{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Nixtla Client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp nixtla_client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide \n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "import logging\n",
    "import math\n",
    "import os\n",
    "import warnings\n",
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union\n",
    "\n",
    "import httpcore\n",
    "import httpx\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import utilsforecast.processing as ufp\n",
    "from fastcore.basics import patch\n",
    "from pydantic import NonNegativeInt, PositiveInt\n",
    "from tenacity import (\n",
    "    RetryCallState,\n",
    "    retry,\n",
    "    retry_if_exception,\n",
    "    retry_if_exception_type,\n",
    "    stop_after_attempt,\n",
    "    stop_after_delay,\n",
    "    wait_fixed,\n",
    ")\n",
    "from utilsforecast.compat import DataFrame, pl_DataFrame\n",
    "from utilsforecast.feature_engineering import _add_time_features, time_features\n",
    "from utilsforecast.validation import ensure_time_dtype, validate_format, validate_freq\n",
    "if TYPE_CHECKING:\n",
    "    import fugue\n",
    "    import triad\n",
    "\n",
    "from nixtla.core.api_error import ApiError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logging.getLogger('httpx').setLevel(logging.ERROR)\n",
    "logger = logging.getLogger(__name__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "from contextlib import contextmanager\n",
    "from itertools import product\n",
    "from time import time, sleep\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "from fastcore.test import test_eq, test_fail, test_warns\n",
    "from utilsforecast.data import generate_series\n",
    "\n",
    "from nixtla.date_features import SpecialDates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "_Loss = Literal[\"default\", \"mae\", \"mse\", \"rmse\", \"mape\", \"smape\"]\n",
    "_Model = Literal[\"azureai\", \"timegpt-1\", \"timegpt-1-long-horizon\"]\n",
    "\n",
    "_date_features_by_freq = {\n",
    "    # Daily frequencies\n",
    "    'B': ['year', 'month', 'day', 'weekday'],\n",
    "    'C': ['year', 'month', 'day', 'weekday'],\n",
    "    'D': ['year', 'month', 'day', 'weekday'],\n",
    "    # Weekly\n",
    "    'W': ['year', 'week', 'weekday'],\n",
    "    # Monthly\n",
    "    'M': ['year', 'month'],\n",
    "    'SM': ['year', 'month', 'day'],\n",
    "    'BM': ['year', 'month'],\n",
    "    'CBM': ['year', 'month'],\n",
    "    'MS': ['year', 'month'],\n",
    "    'SMS': ['year', 'month', 'day'],\n",
    "    'BMS': ['year', 'month'],\n",
    "    'CBMS': ['year', 'month'],\n",
    "    # Quarterly\n",
    "    'Q': ['year', 'quarter'],\n",
    "    'BQ': ['year', 'quarter'],\n",
    "    'QS': ['year', 'quarter'],\n",
    "    'BQS': ['year', 'quarter'],\n",
    "    # Yearly\n",
    "    'A': ['year'],\n",
    "    'Y': ['year'],\n",
    "    'BA': ['year'],\n",
    "    'BY': ['year'],\n",
    "    'AS': ['year'],\n",
    "    'YS': ['year'],\n",
    "    'BAS': ['year'],\n",
    "    'BYS': ['year'],\n",
    "    # Hourly\n",
    "    'BH': ['year', 'month', 'day', 'hour', 'weekday'],\n",
    "    'H': ['year', 'month', 'day', 'hour'],\n",
    "    # Minutely\n",
    "    'T': ['year', 'month', 'day', 'hour', 'minute'],\n",
    "    'min': ['year', 'month', 'day', 'hour', 'minute'],\n",
    "    # Secondly\n",
    "    'S': ['year', 'month', 'day', 'hour', 'minute', 'second'],\n",
    "    # Milliseconds\n",
    "    'L': ['year', 'month', 'day', 'hour', 'minute', 'second', 'millisecond'],\n",
    "    'ms': ['year', 'month', 'day', 'hour', 'minute', 'second', 'millisecond'],\n",
    "    # Microseconds\n",
    "    'U': ['year', 'month', 'day', 'hour', 'minute', 'second', 'microsecond'],\n",
    "    'us': ['year', 'month', 'day', 'hour', 'minute', 'second', 'microsecond'],\n",
    "    # Nanoseconds\n",
    "    'N': []\n",
    "}\n",
    "\n",
    "def _retry_strategy(max_retries: int, retry_interval: int, max_wait_time: int):\n",
    "    def should_retry(exc: Exception) -> bool:\n",
    "        retriable_exceptions = (\n",
    "            ConnectionResetError,\n",
    "            httpcore.ConnectError,\n",
    "            httpcore.RemoteProtocolError,\n",
    "            httpx.ConnectTimeout,\n",
    "            httpx.ReadError,\n",
    "            httpx.RemoteProtocolError,\n",
    "            httpx.ReadTimeout,\n",
    "            httpx.PoolTimeout,\n",
    "            httpx.WriteError,\n",
    "            httpx.WriteTimeout,\n",
    "        )\n",
    "        retriable_codes = [408, 409, 429, 502, 503, 504]\n",
    "        return (\n",
    "            isinstance(exc, retriable_exceptions)\n",
    "            or (isinstance(exc, ApiError) and exc.status_code in retriable_codes)\n",
    "        )\n",
    "\n",
    "    def after_retry(retry_state: RetryCallState) -> None:\n",
    "        error = retry_state.outcome.exception()\n",
    "        logger.error(\n",
    "            f\"Attempt {retry_state.attempt_number} failed with error: {error}\"\n",
    "        )\n",
    "\n",
    "    return retry(\n",
    "        retry=retry_if_exception(should_retry),\n",
    "        wait=wait_fixed(retry_interval),\n",
    "        after=after_retry,\n",
    "        stop=stop_after_attempt(max_retries) | stop_after_delay(max_wait_time),\n",
    "        reraise=True,\n",
    "    )\n",
    "\n",
    "def _maybe_infer_freq(\n",
    "    df: DataFrame,\n",
    "    freq: Optional[str],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    ") -> str:\n",
    "    if freq is not None and freq not in ['W', 'M', 'Q', 'Y', 'A']:\n",
    "        return freq\n",
    "    if isinstance(df, pl_DataFrame):\n",
    "        raise ValueError(\n",
    "            \"Cannot infer frequency for a polars DataFrame, please set the \"\n",
    "            \"`freq` argument to a valid polars offset.\\nYou can find them at \"\n",
    "            \"https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.Expr.dt.offset_by.html\"\n",
    "        )\n",
    "    assert isinstance(df, pd.DataFrame)\n",
    "    sizes = df[id_col].value_counts(sort=True)\n",
    "    times = df.loc[df[id_col] == sizes.index[0], time_col].sort_values()\n",
    "    if times.dt.tz is not None:\n",
    "        times = times.dt.tz_convert('UTC').dt.tz_localize(None)\n",
    "    inferred_freq = pd.infer_freq(times.values)\n",
    "    if inferred_freq is None:\n",
    "        raise RuntimeError(\n",
    "            'Could not infer the frequency of the time column. This could be due '\n",
    "            'to inconsistent intervals. Please check your data for missing, '\n",
    "            'duplicated or irregular timestamps'\n",
    "        )\n",
    "    if freq is not None:\n",
    "        # check we have the same base frequency\n",
    "        # except when we have yearly frequency (A, and Y means the same)\n",
    "        if (\n",
    "            (freq[0] != inferred_freq[0] and freq[0] not in ('A', 'Y'))\n",
    "            or (freq[0] in ('A', 'Y') and inferred_freq[0] not in ('A', 'Y'))\n",
    "        ):\n",
    "            raise RuntimeError(f'Failed to infer special date, inferred freq {inferred_freq}')\n",
    "    logger.info(f'Inferred freq: {inferred_freq}')\n",
    "    return inferred_freq\n",
    "\n",
    "def _standardize_freq(freq: str) -> str:\n",
    "    return freq.replace('mo', 'MS')\n",
    "\n",
    "def _array_tails(\n",
    "    x: np.ndarray,\n",
    "    indptr: np.ndarray,\n",
    "    out_sizes: np.ndarray,\n",
    ") -> np.ndarray:\n",
    "    if (out_sizes > np.diff(indptr)).any():\n",
    "        raise ValueError('out_sizes must be at most the original sizes.')\n",
    "    idxs = np.hstack(\n",
    "        [\n",
    "            np.arange(end - size, end)\n",
    "            for end, size in zip(indptr[1:], out_sizes)\n",
    "        ]\n",
    "    )\n",
    "    return x[idxs]\n",
    "\n",
    "def _tail(proc: ufp.ProcessedDF, n: int) -> ufp.ProcessedDF:\n",
    "    new_sizes = np.minimum(np.diff(proc.indptr), n)\n",
    "    new_indptr = np.append(0, new_sizes.cumsum())\n",
    "    new_data = _array_tails(proc.data, proc.indptr, new_sizes)\n",
    "    return ufp.ProcessedDF(\n",
    "        uids=proc.uids,\n",
    "        last_times=proc.last_times,\n",
    "        data=new_data,\n",
    "        indptr=new_indptr,\n",
    "        sort_idxs=None,\n",
    "    )\n",
    "\n",
    "def _partition_series(\n",
    "    payload: Dict[str, Any], n_part: int, h: int\n",
    ") -> List[Dict[str, Any]]:\n",
    "    parts = []\n",
    "    series = payload.pop(\"series\")\n",
    "    n_series = len(series[\"sizes\"])\n",
    "    n_part = min(n_part, n_series)\n",
    "    series_per_part = math.ceil(n_series / n_part)\n",
    "    prev_size = 0\n",
    "    for i in range(0, n_series, series_per_part):\n",
    "        sizes = series[\"sizes\"][i : i + series_per_part]\n",
    "        curr_size = sum(sizes)\n",
    "        part_idxs = slice(prev_size, prev_size + curr_size)\n",
    "        prev_size += curr_size\n",
    "        part_series = {\n",
    "            \"y\": series[\"y\"][part_idxs],\n",
    "            \"sizes\": sizes,\n",
    "        }\n",
    "        if series[\"X\"] is None:\n",
    "            part_series[\"X\"] = None\n",
    "            if h > 0:\n",
    "                part_series[\"X_future\"] = None\n",
    "        else:\n",
    "            part_series[\"X\"] = [x[part_idxs] for x in series[\"X\"]]\n",
    "            if h > 0:\n",
    "                part_series[\"X_future\"] = [\n",
    "                    x[i * h : (i + series_per_part) * h] for x in series[\"X_future\"]\n",
    "                ]\n",
    "        parts.append({\"series\": part_series, **payload})\n",
    "    return parts\n",
    "\n",
    "def _maybe_add_date_features(\n",
    "    df: DataFrame,\n",
    "    X_df: Optional[DataFrame],\n",
    "    features: Union[bool, List[Union[str, Callable]]],\n",
    "    one_hot: Union[bool, List[str]],\n",
    "    freq: str,\n",
    "    h: int,\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    ") -> Tuple[DataFrame, Optional[DataFrame]]:\n",
    "    if not features:\n",
    "        return df, X_df\n",
    "    if isinstance(features, list):\n",
    "        date_features = features\n",
    "    else:\n",
    "        date_features = _date_features_by_freq.get(freq, [])\n",
    "        if not date_features:\n",
    "            warnings.warn(\n",
    "                f'Non default date features for {freq} '\n",
    "                'please provide a list of date features'\n",
    "            )\n",
    "    # add features\n",
    "    if X_df is None:\n",
    "        df, X_df = time_features(\n",
    "            df=df,\n",
    "            freq=freq,\n",
    "            features=date_features,\n",
    "            h=h,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "        )\n",
    "    else:\n",
    "        df = _add_time_features(df, features=date_features, time_col=time_col)\n",
    "        X_df = _add_time_features(X_df, features=date_features,time_col=time_col)\n",
    "    # one hot\n",
    "    if isinstance(one_hot, list):\n",
    "        features_one_hot = one_hot\n",
    "    elif one_hot:\n",
    "        features_one_hot = [f for f in date_features if not callable(f)]\n",
    "    else:\n",
    "        features_one_hot = []\n",
    "    if features_one_hot:\n",
    "        X_df = ufp.assign_columns(X_df, target_col, 0)\n",
    "        full_df = ufp.vertical_concat([df, X_df])\n",
    "        if isinstance(full_df, pd.DataFrame):\n",
    "            full_df = pd.get_dummies(\n",
    "                full_df, columns=features_one_hot, dtype='float32'\n",
    "            )\n",
    "        else:\n",
    "            full_df = full_df.to_dummies(columns=features_one_hot)\n",
    "        df = ufp.take_rows(full_df, slice(0, df.shape[0]))\n",
    "        X_df = ufp.take_rows(full_df, slice(df.shape[0], full_df.shape[0]))\n",
    "        X_df = ufp.drop_columns(X_df, target_col)\n",
    "        X_df = ufp.drop_index_if_pandas(X_df)\n",
    "    if h == 0:\n",
    "        # time_features returns an empty df, we use it as None here\n",
    "        X_df = None\n",
    "    return df, X_df\n",
    "\n",
    "def _validate_exog(\n",
    "    df: DataFrame,\n",
    "    X_df: Optional[DataFrame],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    ") -> Tuple[DataFrame, Optional[DataFrame]]:\n",
    "    exogs_df = [c for c in df.columns if c not in (id_col, time_col, target_col)]\n",
    "    if X_df is None:\n",
    "        if exogs_df:\n",
    "            warnings.warn(\n",
    "                f'`df` contains the following exogenous features: {exogs_df}, '\n",
    "                'but `X_df` was not provided. They will be ignored.'\n",
    "            )\n",
    "            df = df[[id_col, time_col, target_col]]\n",
    "        return df, None\n",
    "    exogs_X = [c for c in X_df.columns if c not in (id_col, time_col)]\n",
    "    missing_df = set(exogs_X) - set(exogs_df)\n",
    "    if missing_df:\n",
    "        raise ValueError(\n",
    "            'The following exogenous features are present in `X_df` '\n",
    "            f'but not in `df`: {missing_df}.'\n",
    "        )\n",
    "    missing_X_df = set(exogs_df) - set(exogs_X)\n",
    "    if missing_X_df:\n",
    "        warnings.warn(\n",
    "            'The following exogenous features are present in `df` '\n",
    "            f'but not in `X_df`: {missing_X_df}. They will be ignored'\n",
    "        )\n",
    "        exogs_df = [c for c in exogs_df if c in exogs_X]\n",
    "        df = df[[id_col, time_col, target_col, *exogs_df]]\n",
    "    if exogs_df != exogs_X:\n",
    "        # rearrange columns\n",
    "        X_df = X_df[[id_col, time_col, *exogs_df]]\n",
    "    return df, X_df\n",
    "\n",
    "def _validate_input_size(\n",
    "    df: DataFrame,\n",
    "    id_col: str,\n",
    "    model_input_size: int,\n",
    "    model_horizon: int,\n",
    ") -> None:\n",
    "    min_size = ufp.counts_by_id(df, id_col)['counts'].min()\n",
    "    if min_size < model_input_size + model_horizon:\n",
    "        raise ValueError(\n",
    "            'Your time series data is too short '\n",
    "            'Please make sure that your each serie contains '\n",
    "            f'at least {model_input_size + model_horizon} observations.'\n",
    "        )\n",
    "\n",
    "def _prepare_level_and_quantiles(\n",
    "    level: Optional[List[Union[int, float]]], \n",
    "    quantiles: Optional[List[float]],\n",
    ") -> Tuple[List[Union[int, float]], Optional[List[float]]]:\n",
    "    if level is not None and quantiles is not None:\n",
    "        raise ValueError(\n",
    "            \"You should provide `level` or `quantiles`, but not both.\"\n",
    "        )\n",
    "    if quantiles is None:\n",
    "        return level, quantiles\n",
    "    # we recover level from quantiles\n",
    "    if not all(0 < q < 1 for q in quantiles):\n",
    "        raise ValueError(\"`quantiles` should be floats between 0 and 1.\")\n",
    "    level = [abs(int(100 - 200 * q)) for q in quantiles]\n",
    "    return level, quantiles\n",
    "\n",
    "def _maybe_convert_level_to_quantiles(\n",
    "    df: DataFrame,\n",
    "    quantiles: Optional[List[float]],\n",
    ") -> DataFrame:\n",
    "    if quantiles is None:\n",
    "        return df\n",
    "    out_cols = [c for c in df.columns if '-lo-' not in c and '-hi-' not in c]\n",
    "    df = ufp.copy_if_pandas(df, deep=False)\n",
    "    for q in sorted(quantiles):\n",
    "        if q == 0.5:\n",
    "            col = 'TimeGPT'\n",
    "        else:\n",
    "            lv = int(100 - 200 * q)\n",
    "            hi_or_lo = 'lo' if lv > 0 else 'hi'\n",
    "            lv = abs(lv)\n",
    "            col = f\"TimeGPT-{hi_or_lo}-{lv}\"\n",
    "        q_col = f\"TimeGPT-q-{int(q * 100)}\"\n",
    "        df = ufp.assign_columns(df, q_col, df[col])\n",
    "        out_cols.append(q_col)\n",
    "    return df[out_cols]\n",
    "\n",
    "def _preprocess(\n",
    "    df: DataFrame,\n",
    "    X_df: Optional[DataFrame],\n",
    "    h: int,\n",
    "    freq: str,\n",
    "    date_features: Union[bool, List[Union[str, Callable]]],\n",
    "    date_features_to_one_hot: Union[bool, List[str]],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    ") -> Tuple[ufp.ProcessedDF, Optional[DataFrame], List[str]]:\n",
    "    df, X_df = _maybe_add_date_features(\n",
    "        df=df,\n",
    "        X_df=X_df,\n",
    "        features=date_features,\n",
    "        one_hot=date_features_to_one_hot,\n",
    "        freq=freq,\n",
    "        h=h,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "    )\n",
    "    processed = ufp.process_df(\n",
    "        df=df, id_col=id_col, time_col=time_col, target_col=target_col\n",
    "    )\n",
    "    if X_df is not None:\n",
    "        X_df = ensure_time_dtype(X_df, time_col=time_col)\n",
    "        processed_X = ufp.process_df(\n",
    "            df=X_df, id_col=id_col, time_col=time_col, target_col=None,\n",
    "        )\n",
    "        X_future = processed_X.data.T.tolist()\n",
    "    else:\n",
    "        X_future = None\n",
    "    x_cols = [c for c in df.columns if c not in (id_col, time_col, target_col)]\n",
    "    return processed, X_future, x_cols\n",
    "\n",
    "def _forecast_payload_to_in_sample(payload):\n",
    "    in_sample_payload = {\n",
    "        k: v\n",
    "        for k, v in payload.items()\n",
    "        if k not in ('h', 'finetune_steps', 'finetune_loss')\n",
    "    }\n",
    "    del in_sample_payload['series']['X_future']\n",
    "    return in_sample_payload\n",
    "\n",
    "def _maybe_add_intervals(\n",
    "    df: DataFrame,\n",
    "    intervals: Optional[Dict[str, list[float]]],\n",
    ") -> DataFrame:\n",
    "    if intervals is None:\n",
    "        return df\n",
    "    intervals_df = type(df)(\n",
    "        {f'TimeGPT-{k}': intervals[k] for k in sorted(intervals.keys())}\n",
    "    )\n",
    "    return ufp.horizontal_concat([df, intervals_df])\n",
    "\n",
    "def _maybe_drop_id(df: DataFrame, id_col: str, drop: bool) -> DataFrame:\n",
    "    if drop:\n",
    "        df = ufp.drop_columns(df, id_col)\n",
    "    return df\n",
    "\n",
    "def _parse_in_sample_output(\n",
    "    in_sample_output: Dict[str, Union[list[float], Dict[str, list[float]]]],\n",
    "    df: DataFrame,\n",
    "    processed: ufp.ProcessedDF,\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    ") -> DataFrame:\n",
    "    times = df[time_col].to_numpy()\n",
    "    targets = df[target_col].to_numpy()\n",
    "    if processed.sort_idxs is not None:\n",
    "        times = times[processed.sort_idxs]\n",
    "        targets = targets[processed.sort_idxs]\n",
    "    times = _array_tails(\n",
    "        times, processed.indptr, in_sample_output['sizes']\n",
    "    )\n",
    "    targets = _array_tails(\n",
    "        targets, processed.indptr, in_sample_output['sizes']\n",
    "    )\n",
    "    uids = ufp.repeat(processed.uids, in_sample_output['sizes'])\n",
    "    out = type(df)(\n",
    "        {\n",
    "            id_col: uids,\n",
    "            time_col: times,\n",
    "            target_col: targets,\n",
    "            'TimeGPT': in_sample_output['mean'],\n",
    "        }\n",
    "    )\n",
    "    return _maybe_add_intervals(out, in_sample_output['intervals'])\n",
    "\n",
    "def _restrict_input_samples(level, input_size, model_horizon, h) -> int:\n",
    "    if level is not None:\n",
    "        # add sufficient info to compute\n",
    "        # conformal interval\n",
    "        # @AzulGarza\n",
    "        #  this is an old opinionated decision\n",
    "        #  about reducing the data sent to the api\n",
    "        #  to reduce latency when\n",
    "        #  a user passes level. since currently the model\n",
    "        #  uses conformal prediction, we can change a minimum\n",
    "        #  amount of data if the series are too large\n",
    "        new_input_size = 3 * input_size + max(model_horizon, h)\n",
    "    else:\n",
    "        # we only want to forecast\n",
    "        new_input_size = input_size\n",
    "    return new_input_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| export\n",
    "class NixtlaClient:\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        api_key: Optional[str] = None,\n",
    "        base_url: Optional[str] = None,\n",
    "        timeout: int = 60,\n",
    "        max_retries: int = 6,\n",
    "        retry_interval: int = 10,\n",
    "        max_wait_time: int = 6 * 60,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Client to interact with the Nixtla API.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        api_key : str, optional (default=None)\n",
    "            The authorization api_key interacts with the Nixtla API.\n",
    "            If not provided, will use the NIXTLA_API_KEY environment variable.\n",
    "        base_url : str, optional (default=None)\n",
    "            Custom base_url.\n",
    "            If not provided, will use the NIXTLA_BASE_URL environment variable.\n",
    "        timeout : int, optional (default=60)\n",
    "            Request timeout in seconds. Set this to `None` to disable it.\n",
    "        max_retries : int (default=6)\n",
    "            The maximum number of attempts to make when calling the API before giving up. \n",
    "            It defines how many times the client will retry the API call if it fails. \n",
    "            Default value is 6, indicating the client will attempt the API call up to 6 times in total\n",
    "        retry_interval : int (default=10)\n",
    "            The interval in seconds between consecutive retry attempts. \n",
    "            This is the waiting period before the client tries to call the API again after a failed attempt. \n",
    "            Default value is 10 seconds, meaning the client waits for 10 seconds between retries.\n",
    "        max_wait_time : int (default=360)\n",
    "            The maximum total time in seconds that the client will spend on all retry attempts before giving up. \n",
    "            This sets an upper limit on the cumulative waiting time for all retry attempts. \n",
    "            If this time is exceeded, the client will stop retrying and raise an exception. \n",
    "            Default value is 360 seconds, meaning the client will cease retrying if the total time \n",
    "            spent on retries exceeds 360 seconds. \n",
    "            The client throws a ReadTimeout error after 60 seconds of inactivity. If you want to \n",
    "            catch these errors, use max_wait_time >> 60. \n",
    "        \"\"\"\n",
    "        if api_key is None:\n",
    "            api_key = os.environ['NIXTLA_API_KEY']\n",
    "        if base_url is None:\n",
    "            base_url = os.getenv('NIXTLA_BASE_URL', 'https://api.nixtla.io')\n",
    "        self._client_kwargs = {\n",
    "            'base_url': base_url,\n",
    "            'headers': {'Authorization': f'Bearer {api_key}'},\n",
    "            'timeout': timeout,\n",
    "        }\n",
    "        self._retry_strategy = _retry_strategy(\n",
    "            max_retries=max_retries, retry_interval=retry_interval, max_wait_time=max_wait_time\n",
    "        )\n",
    "        self._model_params: Dict[Tuple[str, str], Tuple[int, int]] = {}\n",
    "\n",
    "    def _make_request(self, client: httpx.Client, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:\n",
    "        resp = client.request(\n",
    "            method='post',\n",
    "            url=endpoint,\n",
    "            json=payload,\n",
    "        )\n",
    "        resp_body = resp.json()\n",
    "        if resp.status_code != 200:\n",
    "            raise ApiError(status_code=resp.status_code, body=resp_body)\n",
    "        if 'data' in resp_body:\n",
    "            resp_body = resp_body['data']\n",
    "        return resp_body\n",
    "\n",
    "    def _make_request_with_retries(\n",
    "        self,\n",
    "        client: httpx.Client,\n",
    "        endpoint: str,\n",
    "        payload: Dict[str, Any],\n",
    "    ) -> Dict[str, Any]:\n",
    "        return self._retry_strategy(self._make_request)(\n",
    "            client=client,\n",
    "            endpoint=endpoint,\n",
    "            payload=payload,\n",
    "        )\n",
    "\n",
    "    def _make_partitioned_requests(\n",
    "        self,\n",
    "        client: httpx.Client,\n",
    "        endpoint: str,\n",
    "        payloads: List[Dict[str, Any]],\n",
    "    ) -> Dict[str, Any]:\n",
    "        from tqdm.auto import tqdm\n",
    "\n",
    "        num_partitions = len(payloads)\n",
    "        results = num_partitions * [None]\n",
    "        max_workers = min(10, num_partitions)\n",
    "        with ThreadPoolExecutor(max_workers) as executor:\n",
    "            future2pos = {\n",
    "                executor.submit(\n",
    "                    self._make_request_with_retries, client, endpoint, payload\n",
    "                ): i\n",
    "                for i, payload in enumerate(payloads)\n",
    "            }\n",
    "            for future in tqdm(as_completed(future2pos), total=len(future2pos)):\n",
    "                pos = future2pos[future]\n",
    "                results[pos] = future.result()\n",
    "        resp = {\"mean\": np.hstack([res[\"mean\"] for res in results])}\n",
    "        for k in ('sizes', 'anomaly', 'y'):\n",
    "            if k in results[0]:\n",
    "                resp[k] = np.hstack([res[k] for res in results])\n",
    "        first_res = results[0]\n",
    "        if first_res[\"intervals\"] is None:\n",
    "            resp[\"intervals\"] = None\n",
    "        else:\n",
    "            resp[\"intervals\"] = {}\n",
    "            for k in first_res[\"intervals\"].keys():\n",
    "                resp[\"intervals\"][k] = np.hstack(\n",
    "                    [res[\"intervals\"][k] for res in results]\n",
    "                )\n",
    "        if \"weights_x\" not in first_res or first_res[\"weights_x\"] is None:\n",
    "            resp[\"weights_x\"] = None\n",
    "        else:\n",
    "            resp[\"weights_x\"] = [res[\"weights_x\"] for res in results]\n",
    "        return resp\n",
    "\n",
    "    def _get_model_params(self, model: str, freq: str) -> Tuple[int, int]:\n",
    "        key = (model, freq)\n",
    "        if key not in self._model_params:\n",
    "            logger.info('Querying model metadata...')\n",
    "            payload = {'model': model, 'freq': freq}\n",
    "            with httpx.Client(**self._client_kwargs) as client:\n",
    "                params = self._make_request_with_retries(\n",
    "                    client, 'model_params', payload\n",
    "                )['detail']\n",
    "            self._model_params[key] = (params['input_size'], params['horizon'])\n",
    "        return self._model_params[key]\n",
    "\n",
    "    def _maybe_assign_weights(\n",
    "        self,\n",
    "        weights: Optional[Union[list[float, list[list[float]]]]],\n",
    "        df: DataFrame,\n",
    "        x_cols: List[str],\n",
    "    ) -> None:\n",
    "        if weights is None:\n",
    "            return\n",
    "        if isinstance(weights[0], list):\n",
    "            self.weights_x = [\n",
    "                type(df)({'features': x_cols, 'weights': w}) for w in weights\n",
    "            ]\n",
    "        else:\n",
    "            self.weights_x = type(df)(\n",
    "                {'features': x_cols, 'weights': weights}\n",
    "            )\n",
    "\n",
    "    def _run_validations(\n",
    "        self,\n",
    "        df: DataFrame,\n",
    "        X_df: Optional[DataFrame],\n",
    "        id_col: str,\n",
    "        time_col: str,\n",
    "        target_col: str,\n",
    "        model: str,\n",
    "        validate_api_key: bool,\n",
    "    ) -> Tuple[DataFrame, Optional[DataFrame], bool]:\n",
    "        if validate_api_key and not self.validate_api_key(log=False):\n",
    "            raise Exception('API Key not valid, please email ops@nixtla.io')\n",
    "        supported_models = ['timegpt-1', 'timegpt-1-long-horizon']\n",
    "        if model not in supported_models:\n",
    "            raise ValueError(\n",
    "                f'unsupported model: {model}. supported models: {supported_models}'\n",
    "            )\n",
    "        drop_id = id_col not in df.columns\n",
    "        if drop_id:\n",
    "            df = ufp.copy_if_pandas(df, deep=False)\n",
    "            df = ufp.assign_columns(df, id_col, 0)\n",
    "            if X_df is not None:\n",
    "                X_df = ufp.copy_if_pandas(X_df, deep=False)\n",
    "                X_df = ufp.assign_columns(X_df, id_col, 0)\n",
    "        if (\n",
    "            isinstance(df, pd.DataFrame)\n",
    "            and time_col not in df\n",
    "            and pd.api.types.is_datetime64_any_dtype(df.index)\n",
    "        ):\n",
    "            df.index.name = time_col\n",
    "            df = df.reset_index()\n",
    "        df = ensure_time_dtype(df, time_col=time_col)\n",
    "        validate_format(df=df, id_col=id_col, time_col=time_col, target_col=target_col)\n",
    "        if ufp.is_nan_or_none(df[target_col]).any():\n",
    "            raise ValueError(f'Target column ({target_col}) cannot contain missing values.')\n",
    "        return df, X_df, drop_id\n",
    "\n",
    "    def validate_api_key(self, log: bool = True) -> bool:\n",
    "        \"\"\"Returns True if your api_key is valid.\"\"\"\n",
    "        try:\n",
    "            with httpx.Client(**self._client_kwargs) as client:\n",
    "                validation = self._make_request_with_retries(\n",
    "                    client, 'validate_token', {}\n",
    "                )\n",
    "        except:\n",
    "            validation = {}\n",
    "        if 'support' in validation and log:\n",
    "            logger.info(f'Happy Forecasting! :), {validation[\"support\"]}')\n",
    "        return (\n",
    "            validation.get('message', '') == 'success'\n",
    "            or 'Forecasting! :)' in validation.get('detail', '')\n",
    "        )\n",
    "\n",
    "    def forecast(\n",
    "        self,\n",
    "        df: DataFrame,\n",
    "        h: PositiveInt,\n",
    "        freq: Optional[str] = None,    \n",
    "        id_col: str = 'unique_id',\n",
    "        time_col: str = 'ds',\n",
    "        target_col: str = 'y',\n",
    "        X_df: Optional[DataFrame] = None,\n",
    "        level: Optional[List[Union[int, float]]] = None,\n",
    "        quantiles: Optional[List[float]] = None,\n",
    "        finetune_steps: NonNegativeInt = 0,\n",
    "        finetune_loss: _Loss = 'default',\n",
    "        clean_ex_first: bool = True,\n",
    "        validate_api_key: bool = False,\n",
    "        add_history: bool = False,\n",
    "        date_features: Union[bool, List[Union[str, Callable]]] = False,\n",
    "        date_features_to_one_hot: Union[bool, List[str]] = False,\n",
    "        model: _Model = 'timegpt-1',\n",
    "        num_partitions: Optional[PositiveInt] = None,\n",
    "    ) -> DataFrame:\n",
    "        \"\"\"Forecast your time series using TimeGPT.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame\n",
    "            The DataFrame on which the function will operate. Expected to contain at least the following columns:\n",
    "            - time_col:\n",
    "                Column name in `df` that contains the time indices of the time series. This is typically a datetime\n",
    "                column with regular intervals, e.g., hourly, daily, monthly data points.\n",
    "            - target_col:\n",
    "                Column name in `df` that contains the target variable of the time series, i.e., the variable we \n",
    "                wish to predict or analyze.\n",
    "            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:\n",
    "            - id_col:\n",
    "                Column name in `df` that identifies unique time series. Each unique value in this column\n",
    "                corresponds to a unique time series.\n",
    "        h : int\n",
    "            Forecast horizon.\n",
    "        freq : str\n",
    "            Frequency of the data. By default, the freq will be inferred automatically.\n",
    "            See [pandas' available frequencies](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases).\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.\n",
    "        X_df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.\n",
    "        level : List[float], optional (default=None)\n",
    "            Confidence levels between 0 and 100 for prediction intervals.\n",
    "        quantiles : List[float], optional (default=None)\n",
    "            Quantiles to forecast, list between (0, 1).\n",
    "            `level` and `quantiles` should not be used simultaneously.\n",
    "            The output dataframe will have the quantile columns\n",
    "            formatted as TimeGPT-q-(100 * q) for each q.\n",
    "            100 * q represents percentiles but we choose this notation\n",
    "            to avoid having dots in column names.\n",
    "        finetune_steps : int (default=0)\n",
    "            Number of steps used to finetune learning TimeGPT in the\n",
    "            new data.\n",
    "        finetune_loss : str (default='default')\n",
    "            Loss function to use for finetuning. Options are: `default`, `mae`, `mse`, `rmse`, `mape`, and `smape`.\n",
    "        clean_ex_first : bool (default=True)\n",
    "            Clean exogenous signal before making forecasts using TimeGPT.\n",
    "        validate_api_key : bool (default=False)\n",
    "            If True, validates api_key before sending requests.\n",
    "        add_history : bool (default=False)\n",
    "            Return fitted values of the model.\n",
    "        date_features : bool or list of str or callable, optional (default=False)\n",
    "            Features computed from the dates. \n",
    "            Can be pandas date attributes or functions that will take the dates as input.\n",
    "            If True automatically adds most used date features for the \n",
    "            frequency of `df`.\n",
    "        date_features_to_one_hot : bool or list of str (default=False)\n",
    "            Apply one-hot encoding to these date features.\n",
    "            If `date_features=True`, then all date features are\n",
    "            one-hot encoded by default.\n",
    "        model : str (default='timegpt-1')\n",
    "            Model to use as a string. Options are: `timegpt-1`, and `timegpt-1-long-horizon`. \n",
    "            We recommend using `timegpt-1-long-horizon` for forecasting \n",
    "            if you want to predict more than one seasonal \n",
    "            period given the frequency of your data.\n",
    "        num_partitions : int (default=None)\n",
    "            Number of partitions to use.\n",
    "            If None, the number of partitions will be equal\n",
    "            to the available parallel resources in distributed environments.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        pandas or polars DataFrame\n",
    "            DataFrame with TimeGPT forecasts for point predictions and probabilistic\n",
    "            predictions (if level is not None).\n",
    "        \"\"\"\n",
    "        if not isinstance(df, (pd.DataFrame, pl_DataFrame)):\n",
    "            return self._distributed_forecast(\n",
    "                df=df,\n",
    "                h=h,\n",
    "                freq=freq,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "                X_df=X_df,\n",
    "                level=level,\n",
    "                quantiles=quantiles,\n",
    "                finetune_steps=finetune_steps,\n",
    "                finetune_loss=finetune_loss,\n",
    "                clean_ex_first=clean_ex_first,\n",
    "                validate_api_key=validate_api_key,\n",
    "                add_history=add_history,\n",
    "                date_features=date_features,\n",
    "                date_features_to_one_hot=date_features_to_one_hot,\n",
    "                model=model,\n",
    "                num_partitions=num_partitions,\n",
    "            )\n",
    "        self.__dict__.pop('weights_x', None)\n",
    "        logger.info('Validating inputs...')\n",
    "        df, X_df, drop_id = self._run_validations(\n",
    "            df=df,\n",
    "            X_df=X_df,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "            validate_api_key=validate_api_key,\n",
    "            model=model,\n",
    "        )\n",
    "        df, X_df = _validate_exog(\n",
    "            df, X_df, id_col=id_col, time_col=time_col, target_col=target_col\n",
    "        )\n",
    "        level, quantiles = _prepare_level_and_quantiles(level, quantiles)\n",
    "        freq = _maybe_infer_freq(df, freq=freq, id_col=id_col, time_col=time_col)\n",
    "        standard_freq = _standardize_freq(freq)\n",
    "        model_input_size, model_horizon = self._get_model_params(model, standard_freq)\n",
    "        if finetune_steps > 0 or level is not None or add_history:\n",
    "            _validate_input_size(df, id_col, model_input_size, model_horizon)\n",
    "        if h > model_horizon:\n",
    "            logger.warning(\n",
    "                'The specified horizon \"h\" exceeds the model horizon. '\n",
    "                'This may lead to less accurate forecasts. '\n",
    "                'Please consider using a smaller horizon.'  \n",
    "            )\n",
    "\n",
    "        logger.info('Preprocessing dataframes...')\n",
    "        processed, X_future, x_cols = _preprocess(\n",
    "            df=df,\n",
    "            X_df=X_df,\n",
    "            h=h,\n",
    "            freq=standard_freq,\n",
    "            date_features=date_features,\n",
    "            date_features_to_one_hot=date_features_to_one_hot,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "        restrict_input = finetune_steps == 0 and not x_cols and not add_history\n",
    "        if restrict_input:\n",
    "            logger.info('Restricting input...')\n",
    "            new_input_size = _restrict_input_samples(\n",
    "                level=level,\n",
    "                input_size=model_input_size,\n",
    "                model_horizon=model_horizon,\n",
    "                h=h,\n",
    "            )\n",
    "            processed = _tail(processed, new_input_size)\n",
    "        if processed.data.shape[1] > 1:\n",
    "            X = processed.data[:, 1:].T.tolist()\n",
    "            logger.info(f'Using the following exogenous features: {x_cols}')\n",
    "        else:\n",
    "            X = None\n",
    "\n",
    "        logger.info('Calling Forecast Endpoint...')\n",
    "        payload = {\n",
    "            'series': {\n",
    "                'y': processed.data[:, 0].tolist(),\n",
    "                'sizes': np.diff(processed.indptr).tolist(),\n",
    "                'X': X,\n",
    "                'X_future': X_future,\n",
    "            },\n",
    "            'model': model,\n",
    "            'h': h,\n",
    "            'freq': standard_freq,\n",
    "            'clean_ex_first': clean_ex_first,\n",
    "            'level': level,\n",
    "            'finetune_steps': finetune_steps,\n",
    "            'finetune_loss': finetune_loss,\n",
    "        }\n",
    "        with httpx.Client(**self._client_kwargs) as client:\n",
    "            if num_partitions is None:\n",
    "                resp = self._make_request_with_retries(client, 'v2/forecast', payload)\n",
    "                if add_history:\n",
    "                    in_sample_payload = _forecast_payload_to_in_sample(payload)\n",
    "                    logger.info('Calling Historical Forecast Endpoint...')\n",
    "                    in_sample_resp = self._make_request_with_retries(\n",
    "                        client, 'v2/historic_forecast', in_sample_payload,\n",
    "                    )\n",
    "            else:\n",
    "                payloads = _partition_series(payload, num_partitions, h)\n",
    "                resp = self._make_partitioned_requests(client, 'v2/forecast', payloads)\n",
    "                if add_history:\n",
    "                    in_sample_payloads = [\n",
    "                        _forecast_payload_to_in_sample(p) for p in payloads\n",
    "                    ]\n",
    "                    logger.info('Calling Historical Forecast Endpoint...')\n",
    "                    in_sample_resp = self._make_partitioned_requests(\n",
    "                        client, 'v2/historic_forecast', in_sample_payloads,\n",
    "                    )\n",
    "\n",
    "        # assemble result\n",
    "        out = ufp.make_future_dataframe(\n",
    "            uids=processed.uids,\n",
    "            last_times=type(processed.uids)(processed.last_times),\n",
    "            freq=freq,\n",
    "            h=h,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "        )\n",
    "        out = ufp.assign_columns(out, 'TimeGPT', resp['mean'])\n",
    "        out = _maybe_add_intervals(out, resp['intervals'])\n",
    "        out = _maybe_convert_level_to_quantiles(out, quantiles)\n",
    "        if add_history:\n",
    "            in_sample_df = _parse_in_sample_output(\n",
    "                in_sample_output=in_sample_resp,\n",
    "                df=df,\n",
    "                processed=processed,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "            )\n",
    "            in_sample_df = ufp.drop_columns(in_sample_df, target_col)\n",
    "            out = ufp.vertical_concat([in_sample_df, out])\n",
    "            out = ufp.sort(out, by=[id_col, time_col])\n",
    "        out = _maybe_drop_id(df=out, id_col=id_col, drop=drop_id)\n",
    "        self._maybe_assign_weights(weights=resp['weights_x'], df=df, x_cols=x_cols)\n",
    "        return out\n",
    "\n",
    "    def detect_anomalies(\n",
    "        self,\n",
    "        df: DataFrame,\n",
    "        freq: Optional[str] = None,    \n",
    "        id_col: str = 'unique_id',\n",
    "        time_col: str = 'ds',\n",
    "        target_col: str = 'y',\n",
    "        level: Union[int, float] = 99,\n",
    "        clean_ex_first: bool = True,\n",
    "        validate_api_key: bool = False,\n",
    "        date_features: Union[bool, List[str]] = False,\n",
    "        date_features_to_one_hot: Union[bool, List[str]] = False,\n",
    "        model: _Model = 'timegpt-1',\n",
    "        num_partitions: Optional[PositiveInt] = None,\n",
    "    ):\n",
    "        \"\"\"Detect anomalies in your time series using TimeGPT.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame\n",
    "            The DataFrame on which the function will operate. Expected to contain at least the following columns:\n",
    "            - time_col:\n",
    "                Column name in `df` that contains the time indices of the time series. This is typically a datetime\n",
    "                column with regular intervals, e.g., hourly, daily, monthly data points.\n",
    "            - target_col:\n",
    "                Column name in `df` that contains the target variable of the time series, i.e., the variable we \n",
    "                wish to predict or analyze.\n",
    "            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:\n",
    "            - id_col:\n",
    "                Column name in `df` that identifies unique time series. Each unique value in this column\n",
    "                corresponds to a unique time series.\n",
    "        freq : str\n",
    "            Frequency of the data. By default, the freq will be inferred automatically.\n",
    "            See [pandas' available frequencies](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases).\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.\n",
    "        level : float (default=99)\n",
    "            Confidence level between 0 and 100 for detecting the anomalies.\n",
    "        clean_ex_first : bool (default=True)\n",
    "            Clean exogenous signal before making forecasts\n",
    "            using TimeGPT.\n",
    "        validate_api_key : bool (default=False)\n",
    "            If True, validates api_key before sending requests.\n",
    "        date_features : bool or list of str or callable, optional (default=False)\n",
    "            Features computed from the dates. \n",
    "            Can be pandas date attributes or functions that will take the dates as input.\n",
    "            If True automatically adds most used date features for the \n",
    "            frequency of `df`.\n",
    "        date_features_to_one_hot : bool or list of str (default=False)\n",
    "            Apply one-hot encoding to these date features.\n",
    "            If `date_features=True`, then all date features are\n",
    "            one-hot encoded by default.\n",
    "        model : str (default='timegpt-1')\n",
    "            Model to use as a string. Options are: `timegpt-1`, and `timegpt-1-long-horizon`. \n",
    "            We recommend using `timegpt-1-long-horizon` for forecasting \n",
    "            if you want to predict more than one seasonal \n",
    "            period given the frequency of your data.\n",
    "        num_partitions : int (default=None)\n",
    "            Number of partitions to use.\n",
    "            If None, the number of partitions will be equal\n",
    "            to the available parallel resources in distributed environments.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        pandas or polars DataFrame\n",
    "            DataFrame with anomalies flagged by TimeGPT.\n",
    "        \"\"\"\n",
    "        if not isinstance(df, (pd.DataFrame, pl_DataFrame)):\n",
    "            return self._distributed_detect_anomalies(\n",
    "                df=df,\n",
    "                freq=freq,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "                level=level,\n",
    "                clean_ex_first=clean_ex_first,\n",
    "                validate_api_key=validate_api_key,\n",
    "                date_features=date_features,\n",
    "                date_features_to_one_hot=date_features_to_one_hot,\n",
    "                model=model,\n",
    "                num_partitions=num_partitions,\n",
    "            )\n",
    "        self.__dict__.pop('weights_x', None)\n",
    "        df, _, drop_id = self._run_validations(\n",
    "            df=df,\n",
    "            X_df=None,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "            validate_api_key=validate_api_key,\n",
    "            model=model,\n",
    "        )\n",
    "        freq = _maybe_infer_freq(df, freq=freq, id_col=id_col, time_col=time_col)\n",
    "        standard_freq = _standardize_freq(freq)\n",
    "        model_input_size, model_horizon = self._get_model_params(model, standard_freq)\n",
    "\n",
    "        logger.info('Preprocessing dataframes...')\n",
    "        processed, _, x_cols = _preprocess(\n",
    "            df=df,\n",
    "            X_df=None,\n",
    "            h=0,\n",
    "            freq=standard_freq,\n",
    "            date_features=date_features,\n",
    "            date_features_to_one_hot=date_features_to_one_hot,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "        if processed.data.shape[1] > 1:\n",
    "            X = processed.data[:, 1:].T.tolist()\n",
    "            logger.info(f'Using the following exogenous features: {x_cols}')\n",
    "        else:\n",
    "            X = None\n",
    "\n",
    "        logger.info('Calling Anomaly Detector Endpoint...')\n",
    "        payload = {\n",
    "            'series': {\n",
    "                'y': processed.data[:, 0].tolist(),\n",
    "                'sizes': np.diff(processed.indptr).tolist(),\n",
    "                'X': X,\n",
    "            },\n",
    "            'model': model,\n",
    "            'freq': standard_freq,\n",
    "            'clean_ex_first': clean_ex_first,\n",
    "            'level': level,\n",
    "        }\n",
    "        with httpx.Client(**self._client_kwargs) as client:\n",
    "            if num_partitions is None:\n",
    "                resp = self._make_request_with_retries(\n",
    "                    client, 'v2/anomaly_detection', payload\n",
    "                )\n",
    "            else:\n",
    "                payloads = _partition_series(payload, num_partitions, h=0)\n",
    "                resp = self._make_partitioned_requests(client, 'v2/anomaly_detection', payloads)\n",
    "\n",
    "        # assemble result\n",
    "        out = _parse_in_sample_output(\n",
    "            in_sample_output=resp,\n",
    "            df=df,\n",
    "            processed=processed,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "        out = ufp.assign_columns(out, 'anomaly', resp['anomaly'])\n",
    "        out = _maybe_drop_id(df=out, id_col=id_col, drop=drop_id)\n",
    "        self._maybe_assign_weights(weights=resp['weights_x'], df=df, x_cols=x_cols)\n",
    "        return out\n",
    "\n",
    "    def cross_validation(\n",
    "        self,\n",
    "        df: DataFrame,\n",
    "        h: PositiveInt,\n",
    "        freq: Optional[str] = None,\n",
    "        id_col: str = \"unique_id\",\n",
    "        time_col: str = \"ds\",\n",
    "        target_col: str = \"y\",\n",
    "        level: Optional[List[Union[int, float]]] = None,\n",
    "        quantiles: Optional[List[float]] = None,\n",
    "        validate_api_key: bool = False,\n",
    "        n_windows: PositiveInt = 1,\n",
    "        step_size: Optional[PositiveInt] = None,\n",
    "        finetune_steps: NonNegativeInt = 0,\n",
    "        finetune_loss: str = 'default',\n",
    "        clean_ex_first: bool = True,\n",
    "        date_features: Union[bool, List[str]] = False,\n",
    "        date_features_to_one_hot: Union[bool, List[str]] = False,\n",
    "        model: str = 'timegpt-1',\n",
    "        num_partitions: Optional[PositiveInt] = None,\n",
    "    ):\n",
    "        \"\"\"Perform cross validation in your time series using TimeGPT.\n",
    "        \n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame\n",
    "            The DataFrame on which the function will operate. Expected to contain at least the following columns:\n",
    "            - time_col:\n",
    "                Column name in `df` that contains the time indices of the time series. This is typically a datetime\n",
    "                column with regular intervals, e.g., hourly, daily, monthly data points.\n",
    "            - target_col:\n",
    "                Column name in `df` that contains the target variable of the time series, i.e., the variable we\n",
    "                wish to predict or analyze.\n",
    "            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:\n",
    "            - id_col:\n",
    "                Column name in `df` that identifies unique time series. Each unique value in this column\n",
    "                corresponds to a unique time series.\n",
    "        h : int\n",
    "            Forecast horizon.\n",
    "        freq : str\n",
    "            Frequency of the data. By default, the freq will be inferred automatically.\n",
    "            See [pandas' available frequencies](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases).\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.\n",
    "        level : float (default=99)\n",
    "            Confidence level between 0 and 100 for prediction intervals.\n",
    "        quantiles : List[float], optional (default=None)\n",
    "            Quantiles to forecast, list between (0, 1).\n",
    "            `level` and `quantiles` should not be used simultaneously.\n",
    "            The output dataframe will have the quantile columns\n",
    "            formatted as TimeGPT-q-(100 * q) for each q.\n",
    "            100 * q represents percentiles but we choose this notation\n",
    "            to avoid having dots in column names.\n",
    "        validate_api_key : bool (default=False)\n",
    "            If True, validates api_key before sending requests.\n",
    "        n_windows : int (defaul=1)\n",
    "            Number of windows to evaluate.\n",
    "        step_size : int, optional (default=None)\n",
    "            Step size between each cross validation window. If None it will be equal to `h`.\n",
    "        finetune_steps : int (default=0)\n",
    "            Number of steps used to finetune TimeGPT in the\n",
    "            new data.\n",
    "        finetune_loss : str (default='default')\n",
    "            Loss function to use for finetuning. Options are: `default`, `mae`, `mse`, `rmse`, `mape`, and `smape`.\n",
    "        clean_ex_first : bool (default=True)\n",
    "            Clean exogenous signal before making forecasts\n",
    "            using TimeGPT.\n",
    "        date_features : bool or list of str or callable, optional (default=False)\n",
    "            Features computed from the dates.\n",
    "            Can be pandas date attributes or functions that will take the dates as input.\n",
    "            If True automatically adds most used date features for the\n",
    "            frequency of `df`.\n",
    "        date_features_to_one_hot : bool or list of str (default=False)\n",
    "            Apply one-hot encoding to these date features.\n",
    "            If `date_features=True`, then all date features are\n",
    "            one-hot encoded by default.\n",
    "        model : str (default='timegpt-1')\n",
    "            Model to use as a string. Options are: `timegpt-1`, and `timegpt-1-long-horizon`. \n",
    "            We recommend using `timegpt-1-long-horizon` for forecasting \n",
    "            if you want to predict more than one seasonal \n",
    "            period given the frequency of your data.\n",
    "        num_partitions : int (default=None)\n",
    "            Number of partitions to use.\n",
    "            If None, the number of partitions will be equal\n",
    "            to the available parallel resources in distributed environments.\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        pandas or polars DataFrame\n",
    "            DataFrame with cross validation forecasts.\n",
    "        \"\"\"\n",
    "        if not isinstance(df, (pd.DataFrame, pl_DataFrame)):\n",
    "            return self._distributed_cross_validation(\n",
    "                df=df,\n",
    "                h=h,\n",
    "                freq=freq,\n",
    "                id_col=id_col,\n",
    "                time_col=time_col,\n",
    "                target_col=target_col,\n",
    "                level=level,\n",
    "                quantiles=quantiles,\n",
    "                n_windows=n_windows,\n",
    "                step_size=step_size,\n",
    "                validate_api_key=validate_api_key,\n",
    "                finetune_steps=finetune_steps,\n",
    "                finetune_loss=finetune_loss,\n",
    "                clean_ex_first=clean_ex_first,\n",
    "                date_features=date_features,\n",
    "                date_features_to_one_hot=date_features_to_one_hot,\n",
    "                model=model,\n",
    "                num_partitions=num_partitions,\n",
    "            )\n",
    "        df, _, drop_id = self._run_validations(\n",
    "            df=df,\n",
    "            X_df=None,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "            validate_api_key=validate_api_key,\n",
    "            model=model,\n",
    "        )\n",
    "        freq = _maybe_infer_freq(df, freq=freq, id_col=id_col, time_col=time_col)\n",
    "        standard_freq = _standardize_freq(freq)\n",
    "        level, quantiles = _prepare_level_and_quantiles(level, quantiles)\n",
    "        model_input_size, model_horizon = self._get_model_params(model, standard_freq)\n",
    "        if step_size is None:\n",
    "            step_size = h\n",
    "\n",
    "        logger.info('Preprocessing dataframes...')\n",
    "        processed, _, x_cols = _preprocess(\n",
    "            df=df,\n",
    "            X_df=None,\n",
    "            h=0,\n",
    "            freq=standard_freq,\n",
    "            date_features=date_features,\n",
    "            date_features_to_one_hot=date_features_to_one_hot,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )\n",
    "        times = df[time_col].to_numpy()\n",
    "        if processed.sort_idxs is not None:\n",
    "            times = times[processed.sort_idxs]\n",
    "        restrict_input = finetune_steps == 0 and not x_cols\n",
    "        if restrict_input:\n",
    "            logger.info('Restricting input...')\n",
    "            new_input_size = _restrict_input_samples(\n",
    "                level=level,\n",
    "                input_size=model_input_size,\n",
    "                model_horizon=model_horizon,\n",
    "                h=h,\n",
    "            )\n",
    "            new_input_size += h + step_size * (n_windows - 1)\n",
    "            orig_indptr = processed.indptr\n",
    "            processed = _tail(processed, new_input_size)\n",
    "            times = _array_tails(times, orig_indptr, np.diff(processed.indptr))\n",
    "        if processed.data.shape[1] > 1:\n",
    "            X = processed.data[:, 1:].T.tolist()\n",
    "            logger.info(f'Using the following exogenous features: {x_cols}')\n",
    "        else:\n",
    "            X = None\n",
    "\n",
    "        logger.info('Calling Cross Validation Endpoint...')\n",
    "        payload = {\n",
    "            'series': {\n",
    "                'y': processed.data[:, 0].tolist(),\n",
    "                'sizes': np.diff(processed.indptr).tolist(),\n",
    "                'X': X,\n",
    "            },\n",
    "            'model': model,\n",
    "            'h': h,\n",
    "            'n_windows': n_windows,\n",
    "            'step_size': step_size,\n",
    "            'freq': standard_freq,\n",
    "            'clean_ex_first': clean_ex_first,\n",
    "            'level': level,\n",
    "            'finetune_steps': finetune_steps,\n",
    "            'finetune_loss': finetune_loss,\n",
    "        }\n",
    "        with httpx.Client(**self._client_kwargs) as client:\n",
    "            if num_partitions is None:\n",
    "                resp = self._make_request_with_retries(\n",
    "                    client, 'v2/cross_validation', payload\n",
    "                )\n",
    "            else:\n",
    "                payloads = _partition_series(payload, num_partitions, h=0)\n",
    "                resp = self._make_partitioned_requests(client, 'v2/cross_validation', payloads)\n",
    "\n",
    "        # assemble result\n",
    "        out = ufp.cv_times(\n",
    "            times=times,\n",
    "            uids=processed.uids,\n",
    "            indptr=processed.indptr,\n",
    "            h=h,\n",
    "            test_size=h + step_size * (n_windows - 1),\n",
    "            step_size=step_size,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "        )\n",
    "        out = ufp.sort(out, by=[id_col, 'cutoff', time_col])\n",
    "        out = ufp.assign_columns(out, target_col, resp['y'])\n",
    "        out = ufp.assign_columns(out, 'TimeGPT', resp['mean'])\n",
    "        out = _maybe_add_intervals(out, resp['intervals'])\n",
    "        out = _maybe_drop_id(df=out, id_col=id_col, drop=drop_id)\n",
    "        return _maybe_convert_level_to_quantiles(out, quantiles)\n",
    "\n",
    "    def plot(\n",
    "        self,\n",
    "        df: Optional[DataFrame] = None,\n",
    "        forecasts_df: Optional[DataFrame] = None,\n",
    "        id_col: str = 'unique_id',\n",
    "        time_col: str = 'ds',\n",
    "        target_col: str = 'y',\n",
    "        unique_ids: Union[Optional[List[str]], np.ndarray] = None,\n",
    "        plot_random: bool = True,\n",
    "        models: Optional[List[str]] = None,\n",
    "        level: Optional[List[float]] = None,\n",
    "        max_insample_length: Optional[int] = None,\n",
    "        plot_anomalies: bool = False,\n",
    "        engine: str = 'matplotlib',\n",
    "        resampler_kwargs: Optional[Dict] = None,\n",
    "    ):\n",
    "        \"\"\"Plot forecasts and insample values.\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        df : pandas or polars DataFrame, optional (default=None)\n",
    "            The DataFrame on which the function will operate. Expected to contain at least the following columns:\n",
    "            - time_col:\n",
    "                Column name in `df` that contains the time indices of the time series. This is typically a datetime\n",
    "                column with regular intervals, e.g., hourly, daily, monthly data points.\n",
    "            - target_col:\n",
    "                Column name in `df` that contains the target variable of the time series, i.e., the variable we \n",
    "                wish to predict or analyze.\n",
    "            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:\n",
    "            - id_col:\n",
    "                Column name in `df` that identifies unique time series. Each unique value in this column\n",
    "                corresponds to a unique time series.\n",
    "        forecasts_df : pandas or polars DataFrame, optional (default=None)\n",
    "            DataFrame with columns [`unique_id`, `ds`] and models.\n",
    "        id_col : str (default='unique_id')\n",
    "            Column that identifies each serie.\n",
    "        time_col : str (default='ds')\n",
    "            Column that identifies each timestep, its values can be timestamps or integers.\n",
    "        target_col : str (default='y')\n",
    "            Column that contains the target.\n",
    "        unique_ids : List[str], optional (default=None)\n",
    "            Time Series to plot.\n",
    "            If None, time series are selected randomly.\n",
    "        plot_random : bool (default=True)\n",
    "            Select time series to plot randomly.\n",
    "        models : List[str], optional (default=None)\n",
    "            List of models to plot.\n",
    "        level : List[float], optional (default=None)\n",
    "            List of prediction intervals to plot if paseed.\n",
    "        max_insample_length : int, optional (default=None)\n",
    "            Max number of train/insample observations to be plotted.\n",
    "        plot_anomalies : bool (default=False)\n",
    "            Plot anomalies for each prediction interval.\n",
    "        engine : str (default='plotly')\n",
    "            Library used to plot. 'plotly', 'plotly-resampler' or 'matplotlib'.\n",
    "        resampler_kwargs : dict\n",
    "            Kwargs to be passed to plotly-resampler constructor.\n",
    "            For further custumization (\"show_dash\") call the method,\n",
    "            store the plotting object and add the extra arguments to\n",
    "            its `show_dash` method.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            from utilsforecast.plotting import plot_series\n",
    "        except ModuleNotFoundError:\n",
    "            raise Exception(\n",
    "                'You have to install additional dependencies to use this method, '\n",
    "                'please install them using `pip install \"nixtla[plotting]\"`'\n",
    "            )\n",
    "        if df is not None and id_col not in df.columns:\n",
    "            df = ufp.copy_if_pandas(df, deep=False)\n",
    "            df = ufp.assign_columns(df, id_col, 'ts_0')\n",
    "        df = ensure_time_dtype(df, time_col=time_col)\n",
    "        if forecasts_df is not None:\n",
    "            if id_col not in forecasts_df.columns:\n",
    "                forecasts_df = ufp.copy_if_pandas(forecasts_df, deep=False)\n",
    "                forecasts_df = ufp.assign_columns(forecasts_df, id_col, 'ts_0')\n",
    "            forecasts_df = ensure_time_dtype(forecasts_df, time_col=time_col)\n",
    "            if 'anomaly' in forecasts_df.columns:\n",
    "                # special case to plot outputs\n",
    "                # from detect_anomalies\n",
    "                df = None\n",
    "                forecasts_df = ufp.drop_columns(forecasts_df, 'anomaly')\n",
    "                cols = [c for c in forecasts_df.columns if 'TimeGPT-lo-' in c]\n",
    "                level = [c.replace('TimeGPT-lo-', '') for c in cols][0]\n",
    "                level = float(level) if '.' in level else int(level)\n",
    "                level = [level]\n",
    "                plot_anomalies = True\n",
    "                models = ['TimeGPT']\n",
    "        return plot_series(\n",
    "            df=df,\n",
    "            forecasts_df=forecasts_df,\n",
    "            ids=unique_ids,\n",
    "            plot_random=plot_random,\n",
    "            models=models,\n",
    "            level=level,\n",
    "            max_insample_length=max_insample_length,\n",
    "            plot_anomalies=plot_anomalies,\n",
    "            engine=engine,\n",
    "            resampler_kwargs=resampler_kwargs,\n",
    "            palette=\"tab20b\",\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "@contextmanager\n",
    "def delete_env_var(key):\n",
    "    original_value = os.environ.get(key)\n",
    "    rm = False\n",
    "    if key in os.environ:\n",
    "        del os.environ[key]\n",
    "        rm = True\n",
    "    try:\n",
    "        yield\n",
    "    finally:\n",
    "        if rm:\n",
    "            os.environ[key] = original_value\n",
    "# test api_key fail\n",
    "with delete_env_var('NIXTLA_API_KEY'), delete_env_var('TIMEGPT_TOKEN'):\n",
    "    test_fail(\n",
    "        lambda: NixtlaClient(),\n",
    "        contains='NIXTLA_API_KEY',\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "nixtla_client = NixtlaClient()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "nixtla_client.validate_api_key()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "_nixtla_client = NixtlaClient(api_key=\"invalid\")\n",
    "test_eq(_nixtla_client.validate_api_key(), False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "_nixtla_client = NixtlaClient(\n",
    "    api_key=os.environ['NIXTLA_API_KEY_CUSTOM'], \n",
    "    base_url=os.environ['NIXTLA_BASE_URL_CUSTOM'],\n",
    ")\n",
    "_nixtla_client.validate_api_key()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "test_fail(\n",
    "    lambda: NixtlaClient(api_key='transphobic').forecast(df=pd.DataFrame(), h=None, validate_api_key=True),\n",
    "    contains='nixtla'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test input_size\n",
    "test_eq(\n",
    "    nixtla_client._get_model_params(model='timegpt-1', freq='D'),\n",
    "    (28, 7),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can start to make forecasts! Let's import an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "df = pd.read_csv(\n",
    "    'https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/air_passengers.csv',\n",
    "    parse_dates=['timestamp'],\n",
    ")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test date_features with multiple series\n",
    "# and different ends\n",
    "test_series = generate_series(n_series=2, min_length=5, max_length=20)\n",
    "h = 12\n",
    "fcst_test_series = nixtla_client.forecast(test_series, h=12, date_features=['dayofweek'])\n",
    "uids = test_series['unique_id']\n",
    "for uid in uids:\n",
    "    test_eq(\n",
    "        fcst_test_series.query('unique_id == @uid')['ds'].tolist(),\n",
    "        pd.date_range(periods=h + 1, start=test_series.query('unique_id == @uid')['ds'].max())[1:].tolist(),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test quantiles\n",
    "test_fail(\n",
    "    lambda: nixtla_client.forecast(\n",
    "        df=df, \n",
    "        h=12, \n",
    "        time_col='timestamp', \n",
    "        target_col='value', \n",
    "        level=[80], \n",
    "        quantiles=[0.2, 0.3]\n",
    "    ),\n",
    "    contains='not both'\n",
    ")\n",
    "test_qls = list(np.arange(0.1, 1, 0.1))\n",
    "exp_q_cols = [f\"TimeGPT-q-{int(100 * q)}\" for q in test_qls]\n",
    "def test_method_qls(method, **kwargs):\n",
    "    df_qls = method(\n",
    "        df=df, \n",
    "        h=12, \n",
    "        time_col='timestamp', \n",
    "        target_col='value', \n",
    "        quantiles=test_qls,\n",
    "        **kwargs\n",
    "    )\n",
    "    assert all(col in df_qls.columns for col in exp_q_cols)\n",
    "    # test monotonicity of quantiles\n",
    "    df_qls.apply(lambda x: x.is_monotonic_increasing, axis=1).sum() == len(exp_q_cols)\n",
    "test_method_qls(nixtla_client.forecast)\n",
    "test_method_qls(nixtla_client.forecast, add_history=True)\n",
    "test_method_qls(nixtla_client.cross_validation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test num partitions\n",
    "# we need to be sure that we can recover the same results\n",
    "# using a for loop\n",
    "# A: be aware that num partitons can produce different results\n",
    "# when used finetune_steps\n",
    "def test_num_partitions_same_results(method: Callable, num_partitions: int, **kwargs):\n",
    "    res_partitioned = method(**kwargs, num_partitions=num_partitions)\n",
    "    res_no_partitioned = method(**kwargs, num_partitions=1)\n",
    "    sort_by = ['unique_id', 'ds']\n",
    "    if 'cutoff' in res_partitioned:\n",
    "        sort_by.extend(['cutoff'])\n",
    "    pd.testing.assert_frame_equal(\n",
    "        res_partitioned.sort_values(sort_by).reset_index(drop=True), \n",
    "        res_no_partitioned.sort_values(sort_by).reset_index(drop=True),\n",
    "        rtol=1e-2,\n",
    "        atol=1e-2,\n",
    "    )\n",
    "\n",
    "freqs = {'D': 7, 'W-THU': 52, 'Q-DEC': 8, '15T': 4 * 24 * 7}\n",
    "for freq, h in freqs.items():\n",
    "    df_freq = generate_series(\n",
    "        10, \n",
    "        min_length=500 if freq != '15T' else 1_200, \n",
    "        max_length=550 if freq != '15T' else 2_000,\n",
    "    )\n",
    "    #df_freq['y'] = df_freq['y'].astype(np.float32)\n",
    "    df_freq['ds'] = df_freq.groupby('unique_id', observed=True)['ds'].transform(\n",
    "        lambda x: pd.date_range(periods=len(x), freq=freq, end='2023-01-01')\n",
    "    )\n",
    "    min_size = df_freq.groupby('unique_id', observed=True).size().min()\n",
    "    test_num_partitions_same_results(\n",
    "        nixtla_client.detect_anomalies,\n",
    "        level=98,\n",
    "        df=df_freq,\n",
    "        num_partitions=2,\n",
    "    )\n",
    "    test_num_partitions_same_results(\n",
    "        nixtla_client.cross_validation,\n",
    "        h=7,\n",
    "        n_windows=2,\n",
    "        df=df_freq,\n",
    "        num_partitions=2,\n",
    "    )\n",
    "    test_num_partitions_same_results(\n",
    "        nixtla_client.forecast,\n",
    "        df=df_freq,\n",
    "        h=7,\n",
    "        add_history=True,\n",
    "        num_partitions=2,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "def test_retry_behavior(side_effect, max_retries=5, retry_interval=5, max_wait_time=40, should_retry=True, sleep_seconds=5):\n",
    "    mock_nixtla_client = NixtlaClient(\n",
    "        max_retries=max_retries, \n",
    "        retry_interval=retry_interval, \n",
    "        max_wait_time=max_wait_time,\n",
    "    )\n",
    "    mock_nixtla_client._make_request = side_effect\n",
    "    init_time = time()\n",
    "    test_fail(\n",
    "        lambda: mock_nixtla_client.forecast(df=df, h=12, time_col='timestamp', target_col='value'),\n",
    "    )\n",
    "    total_mock_time = time() - init_time\n",
    "    if should_retry:\n",
    "        approx_expected_time = min((max_retries - 1) * retry_interval, max_wait_time)\n",
    "        upper_expected_time = min(max_retries * retry_interval, max_wait_time)\n",
    "        assert total_mock_time >= approx_expected_time, \"It is not retrying as expected\"\n",
    "        # preprocessing time before the first api call should be less than 60 seconds\n",
    "        assert total_mock_time - upper_expected_time - (max_retries - 1) * sleep_seconds <= sleep_seconds\n",
    "    else:\n",
    "        assert total_mock_time <= max_wait_time "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# we want the api to retry in these cases\n",
    "def raise_api_error_with_text(*args, **kwargs):\n",
    "    raise ApiError(\n",
    "        status_code=503, \n",
    "        body=\"\"\"\n",
    "        <html><head>\n",
    "        <meta http-equiv=\"content-type\" content=\"text/html;charset=utf-8\">\n",
    "        <title>503 Server Error</title>\n",
    "        </head>\n",
    "        <body text=#000000 bgcolor=#ffffff>\n",
    "        <h1>Error: Server Error</h1>\n",
    "        <h2>The service you requested is not available at this time.<p>Service error -27.</h2>\n",
    "        <h2></h2>\n",
    "        </body></html>\n",
    "        \"\"\")\n",
    "test_retry_behavior(raise_api_error_with_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# we want the api to not retry in these cases\n",
    "# here A is assuming that the endpoint responds always\n",
    "# with a json\n",
    "def raise_api_error_with_json(*args, **kwargs):\n",
    "    raise ApiError(\n",
    "        status_code=422, \n",
    "        body=dict(detail='Please use numbers'),\n",
    "    )\n",
    "test_retry_behavior(raise_api_error_with_json, should_retry=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test resilience of api calls\n",
    "def raise_read_timeout_error(*args, **kwargs):\n",
    "    sleep_seconds = 5\n",
    "    print(f'raising ReadTimeout error after {sleep_seconds} seconds')\n",
    "    sleep(sleep_seconds)\n",
    "    raise httpx.ReadTimeout('Timed out')\n",
    "\n",
    "def raise_http_error(*args, **kwargs):\n",
    "    print('raising HTTP error')\n",
    "    raise ApiError(status_code=503, body='HTTP error')\n",
    "    \n",
    "combs = [\n",
    "    (2, 5, 30),\n",
    "    (10, 1, 5),\n",
    "]\n",
    "side_effects = [raise_read_timeout_error, raise_http_error]\n",
    "\n",
    "for (max_retries, retry_interval, max_wait_time), side_effect in product(combs, side_effects):\n",
    "    test_retry_behavior(\n",
    "        max_retries=max_retries, \n",
    "        retry_interval=retry_interval, \n",
    "        max_wait_time=max_wait_time, \n",
    "        side_effect=side_effect,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "nixtla_client.plot(df, time_col='timestamp', target_col='value', engine='plotly')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test we recover the same <mean> forecasts\n",
    "# with and without restricting input\n",
    "# (add_history)\n",
    "def test_equal_fcsts_add_history(**kwargs):\n",
    "    fcst_no_rest_df = nixtla_client.forecast(**kwargs, add_history=True)\n",
    "    fcst_no_rest_df = fcst_no_rest_df.groupby('unique_id', observed=True).tail(kwargs['h']).reset_index(drop=True)\n",
    "    fcst_rest_df = nixtla_client.forecast(**kwargs)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        fcst_no_rest_df,\n",
    "        fcst_rest_df,\n",
    "    )\n",
    "    return fcst_rest_df\n",
    "\n",
    "freqs = {'D': 7, 'W-THU': 52, 'Q-DEC': 8, '15T': 4 * 24 * 7}\n",
    "for freq, h in freqs.items():\n",
    "    df_freq = generate_series(\n",
    "        10, \n",
    "        min_length=500 if freq != '15T' else 1_200, \n",
    "        max_length=550 if freq != '15T' else 2_000,\n",
    "    )\n",
    "    df_freq['ds'] = df_freq.groupby('unique_id', observed=True)['ds'].transform(\n",
    "        lambda x: pd.date_range(periods=len(x), freq=freq, end='2023-01-01')\n",
    "    )\n",
    "    kwargs = dict(\n",
    "        df=df_freq,\n",
    "        h=h,\n",
    "    )\n",
    "    fcst_1_df = test_equal_fcsts_add_history(**{**kwargs, 'model': 'timegpt-1'})\n",
    "    fcst_2_df = test_equal_fcsts_add_history(**{**kwargs, 'model': 'timegpt-1-long-horizon'})\n",
    "    test_fail(\n",
    "        lambda: pd.testing.assert_frame_equal(fcst_1_df, fcst_2_df),\n",
    "        contains='(column name=\"TimeGPT\") are different',\n",
    "    )\n",
    "    # add test num_partitions    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#test same results custom url\n",
    "nixtla_client_custom = NixtlaClient(\n",
    "    api_key=os.environ['NIXTLA_API_KEY_CUSTOM'], \n",
    "    base_url=os.environ['NIXTLA_BASE_URL_CUSTOM'],\n",
    ")\n",
    "# forecast method\n",
    "fcst_kwargs = dict(\n",
    "    df=df, \n",
    "    h=12, \n",
    "    level=[90, 95], \n",
    "    add_history=True, \n",
    "    time_col='timestamp', \n",
    "    target_col='value',\n",
    ")\n",
    "fcst_df = nixtla_client.forecast(**fcst_kwargs)\n",
    "fcst_df_custom = nixtla_client_custom.forecast(**fcst_kwargs)\n",
    "pd.testing.assert_frame_equal(\n",
    "    fcst_df,\n",
    "    fcst_df_custom,\n",
    ")\n",
    "# anomalies method\n",
    "anomalies_kwargs = dict(\n",
    "    df=df, \n",
    "    level=99,\n",
    "    time_col='timestamp', \n",
    "    target_col='value',\n",
    ")\n",
    "anomalies_df = nixtla_client.detect_anomalies(**anomalies_kwargs)\n",
    "anomalies_df_custom = nixtla_client_custom.detect_anomalies(**anomalies_kwargs)\n",
    "pd.testing.assert_frame_equal(\n",
    "    anomalies_df,\n",
    "    anomalies_df_custom,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test different results for different models\n",
    "fcst_kwargs['model'] = 'timegpt-1'\n",
    "fcst_timegpt_1 = nixtla_client.forecast(**fcst_kwargs)\n",
    "fcst_kwargs['model'] = 'timegpt-1-long-horizon'\n",
    "fcst_timegpt_long = nixtla_client.forecast(**fcst_kwargs)\n",
    "test_fail(\n",
    "    lambda: pd.testing.assert_frame_equal(fcst_timegpt_1[['TimeGPT']], fcst_timegpt_long[['TimeGPT']]),\n",
    "    contains='(column name=\"TimeGPT\") are different'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test different results for different models\n",
    "# cross validation\n",
    "cv_kwargs = dict(\n",
    "    df=df, \n",
    "    h=12, \n",
    "    level=[90, 95], \n",
    "    time_col='timestamp', \n",
    "    target_col='value',\n",
    ")\n",
    "cv_kwargs['model'] = 'timegpt-1'\n",
    "cv_timegpt_1 = nixtla_client.cross_validation(**cv_kwargs)\n",
    "cv_kwargs['model'] = 'timegpt-1-long-horizon'\n",
    "cv_timegpt_long = nixtla_client.cross_validation(**cv_kwargs)\n",
    "test_fail(\n",
    "    lambda: pd.testing.assert_frame_equal(cv_timegpt_1[['TimeGPT']], cv_timegpt_long[['TimeGPT']]),\n",
    "    contains='(column name=\"TimeGPT\") are different'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test different results for different models\n",
    "# anomalies\n",
    "anomalies_kwargs['model'] = 'timegpt-1'\n",
    "anomalies_timegpt_1 = nixtla_client.detect_anomalies(**anomalies_kwargs)\n",
    "anomalies_kwargs['model'] = 'timegpt-1-long-horizon'\n",
    "anomalies_timegpt_long = nixtla_client.detect_anomalies(**anomalies_kwargs)\n",
    "test_fail(\n",
    "    lambda: pd.testing.assert_frame_equal(anomalies_timegpt_1[['TimeGPT']], anomalies_timegpt_long[['TimeGPT']]),\n",
    "    contains='(column name=\"TimeGPT\") are different'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test unsupported model\n",
    "fcst_kwargs['model'] = 'a-model'\n",
    "test_fail(\n",
    "    lambda: nixtla_client.forecast(**fcst_kwargs),\n",
    "    contains='unsupported model',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test unsupported model\n",
    "anomalies_kwargs['model'] = 'my-awesome-model'\n",
    "test_fail(\n",
    "    lambda: nixtla_client.detect_anomalies(**anomalies_kwargs),\n",
    "    contains='unsupported model',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test add date features\n",
    "df_ = df.rename(columns={'timestamp': 'ds', 'value': 'y'})\n",
    "df_.insert(0, 'unique_id', 'AirPassengers')\n",
    "date_features = ['year', 'month']\n",
    "df_date_features, future_df = _maybe_add_date_features(\n",
    "    df=df_,\n",
    "    X_df=None,\n",
    "    h=12, \n",
    "    freq='MS', \n",
    "    features=date_features,\n",
    "    one_hot=False,\n",
    "    id_col='unique_id',\n",
    "    time_col='ds',\n",
    "    target_col='y',\n",
    ")\n",
    "assert all(col in df_date_features for col in date_features)\n",
    "assert all(col in future_df for col in date_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# cross validation tests\n",
    "df_copy = df_.copy()\n",
    "pd.testing.assert_frame_equal(\n",
    "    df_copy,\n",
    "    df_,\n",
    ")\n",
    "df_test = df_.groupby('unique_id').tail(12)\n",
    "df_train = df_.drop(df_test.index)\n",
    "hyps = [\n",
    "    # finetune steps is unstable due\n",
    "    # to numerical reasons\n",
    "    # dict(finetune_steps=2),\n",
    "    dict(),\n",
    "    dict(clean_ex_first=False),\n",
    "    dict(date_features=['month']),\n",
    "    dict(level=[80, 90]),\n",
    "    #dict(level=[80, 90], finetune_steps=2),\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test exogenous variables cv\n",
    "df_ex_ = df_.copy()\n",
    "df_ex_['exogenous_var'] = df_ex_['y'] + np.random.normal(size=len(df_ex_))\n",
    "x_df_test = df_test.drop(columns='y').merge(df_ex_.drop(columns='y'))\n",
    "for hyp in hyps:\n",
    "    logger.info(f'Hyperparameters: {hyp}')\n",
    "    logger.info('\\n\\nPerforming forecast\\n')\n",
    "    fcst_test = nixtla_client.forecast(\n",
    "        df_train.merge(df_ex_.drop(columns='y')), h=12, X_df=x_df_test, **hyp\n",
    "    )\n",
    "    fcst_test = df_test[['unique_id', 'ds', 'y']].merge(fcst_test).astype({'y': 'float64'})\n",
    "    fcst_test = fcst_test.sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
    "    logger.info('\\n\\nPerforming Cross validation\\n')\n",
    "    fcst_cv = nixtla_client.cross_validation(df_ex_, h=12, **hyp)\n",
    "    fcst_cv = fcst_cv.sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
    "    logger.info('\\n\\nVerify difference\\n')\n",
    "    pd.testing.assert_frame_equal(fcst_test, fcst_cv.drop(columns='cutoff'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test finetune cv\n",
    "finetune_cv = nixtla_client.cross_validation(\n",
    "            df=df_,\n",
    "            h=12,\n",
    "            n_windows=1,\n",
    "            finetune_steps=1\n",
    "        )\n",
    "test_eq(finetune_cv is not None, True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "for hyp in hyps:\n",
    "    fcst_test = nixtla_client.forecast(df_train, h=12, **hyp)\n",
    "    fcst_test = df_test[['unique_id', 'ds', 'y']].merge(fcst_test).astype({'y': 'float64'})\n",
    "    fcst_test = fcst_test.sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
    "    fcst_cv = nixtla_client.cross_validation(df_, h=12, **hyp)\n",
    "    fcst_cv = fcst_cv.sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        fcst_test,\n",
    "        fcst_cv.drop(columns='cutoff'),\n",
    "        rtol=1e-2,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "for hyp in hyps:\n",
    "    fcst_test = nixtla_client.forecast(df_train, h=12, **hyp)\n",
    "    fcst_test.insert(2, 'y', df_test['y'].values)\n",
    "    fcst_test = fcst_test.sort_values(['unique_id', 'ds']).reset_index(drop=True).astype({'y': 'float64'})\n",
    "    fcst_cv = nixtla_client.cross_validation(df_, h=12, **hyp)\n",
    "    fcst_cv = fcst_cv.sort_values(['unique_id', 'ds']).reset_index(drop=True)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        fcst_test,\n",
    "        fcst_cv.drop(columns='cutoff'),\n",
    "        rtol=1e-2,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test add callables\n",
    "date_features = [SpecialDates({'first_dates': ['2021-01-1'], 'second_dates': ['2021-01-01']})]\n",
    "df_daily = df_.copy()\n",
    "df_daily['ds'] = pd.date_range(end='2021-01-01', periods=len(df_daily))\n",
    "df_date_features, future_df = _maybe_add_date_features(\n",
    "    df=df_,\n",
    "    X_df=None,\n",
    "    h=12, \n",
    "    freq='D', \n",
    "    features=date_features,\n",
    "    one_hot=False,\n",
    "    id_col='unique_id',\n",
    "    time_col='ds',\n",
    "    target_col='y',\n",
    ")\n",
    "assert all(col in df_date_features for col in ['first_dates', 'second_dates'])\n",
    "assert all(col in future_df for col in ['first_dates', 'second_dates'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test add date features one hot encoded\n",
    "date_features = ['year', 'month']\n",
    "date_features_to_one_hot = ['month']\n",
    "df_date_features, future_df = _maybe_add_date_features(\n",
    "    df=df_,\n",
    "    X_df=None,\n",
    "    h=12, \n",
    "    freq='D', \n",
    "    features=date_features,\n",
    "    one_hot=date_features_to_one_hot,\n",
    "    id_col='unique_id',\n",
    "    time_col='ds',\n",
    "    target_col='y',\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test pass dataframe with index\n",
    "df_ds_index = df_.set_index('ds')[['unique_id', 'y']]\n",
    "df_ds_index.index = pd.DatetimeIndex(df_ds_index.index)\n",
    "fcst_inferred_df_index = nixtla_client.forecast(df_ds_index, h=10)\n",
    "anom_inferred_df_index = nixtla_client.detect_anomalies(df_ds_index)\n",
    "fcst_inferred_df = nixtla_client.forecast(df_[['ds', 'unique_id', 'y']], h=10)\n",
    "anom_inferred_df = nixtla_client.detect_anomalies(df_[['ds', 'unique_id', 'y']])\n",
    "pd.testing.assert_frame_equal(fcst_inferred_df_index, fcst_inferred_df, atol=1e-3)\n",
    "pd.testing.assert_frame_equal(anom_inferred_df_index, anom_inferred_df, atol=1e-3)\n",
    "df_ds_index = df_ds_index.groupby('unique_id').tail(80)\n",
    "for freq in ['Y', 'W-MON', 'Q-DEC', 'H']:\n",
    "    df_ds_index.index = np.concatenate(\n",
    "        df_ds_index['unique_id'].nunique() * [pd.date_range(end='2023-01-01', periods=80, freq=freq)]\n",
    "    )\n",
    "    df_ds_index.index.name = 'ds'\n",
    "    fcst_inferred_df_index = nixtla_client.forecast(df_ds_index, h=10)\n",
    "    df_test = df_ds_index.reset_index()\n",
    "    fcst_inferred_df = nixtla_client.forecast(df_test, h=10)\n",
    "    pd.testing.assert_frame_equal(fcst_inferred_df_index, fcst_inferred_df, atol=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test add date features with exogenous variables \n",
    "# and multiple series\n",
    "date_features = ['year', 'month']\n",
    "df_actual_future = df_.tail(12)[['unique_id', 'ds']]\n",
    "df_date_features, future_df = _maybe_add_date_features(\n",
    "    df=df_,\n",
    "    X_df=df_actual_future,\n",
    "    h=24, \n",
    "    freq='H', \n",
    "    features=date_features,\n",
    "    one_hot=False,\n",
    "    id_col='unique_id',\n",
    "    time_col='ds',\n",
    "    target_col='y',\n",
    ")\n",
    "assert all(col in df_date_features for col in date_features)\n",
    "assert all(col in future_df for col in date_features)\n",
    "pd.testing.assert_frame_equal(\n",
    "    df_date_features[df_.columns],\n",
    "    df_,\n",
    ")\n",
    "pd.testing.assert_frame_equal(\n",
    "    future_df[df_actual_future.columns],\n",
    "    df_actual_future,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test add date features one hot with exogenous variables \n",
    "# and multiple series\n",
    "date_features = ['month', 'day']\n",
    "df_date_features, future_df = _maybe_add_date_features(\n",
    "    df=df_,\n",
    "    X_df=df_actual_future,\n",
    "    h=24, \n",
    "    freq='H', \n",
    "    features=date_features,\n",
    "    one_hot=date_features,\n",
    "    id_col='unique_id',\n",
    "    time_col='ds',\n",
    "    target_col='y',\n",
    ")\n",
    "pd.testing.assert_frame_equal(\n",
    "    df_date_features[df_.columns],\n",
    "    df_,\n",
    ")\n",
    "pd.testing.assert_frame_equal(\n",
    "    future_df[df_actual_future.columns],\n",
    "    df_actual_future.reset_index(drop=True),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test warning horizon too long\n",
    "nixtla_client.forecast(df=df.tail(3), h=100, time_col='timestamp', target_col='value')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide \n",
    "# test short horizon with add_history\n",
    "test_fail(\n",
    "    lambda: nixtla_client.forecast(df=df.tail(3), h=12, time_col='timestamp', target_col='value', add_history=True),\n",
    "    contains='make sure'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide \n",
    "# test short horizon with finetunning\n",
    "test_fail(\n",
    "    lambda: nixtla_client.forecast(df=df.tail(3), h=12, time_col='timestamp', target_col='value', finetune_steps=10, finetune_loss='mae'),\n",
    "    contains='make sure'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide \n",
    "# test short horizon with level\n",
    "test_fail(\n",
    "    lambda: nixtla_client.forecast(df=df.tail(3), h=12, time_col='timestamp', target_col='value', level=[80, 90]),\n",
    "    contains='make sure'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test custom url\n",
    "# same results\n",
    "_timegpt_fcst_df = _nixtla_client.forecast(df=df, h=12, time_col='timestamp', target_col='value')\n",
    "timegpt_fcst_df = nixtla_client.forecast(df=df, h=12, time_col='timestamp', target_col='value')\n",
    "pd.testing.assert_frame_equal(\n",
    "    _timegpt_fcst_df,\n",
    "    timegpt_fcst_df,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test using index as time_col\n",
    "# same results\n",
    "df_test = df.copy()\n",
    "df_test[\"timestamp\"] = pd.to_datetime(df_test[\"timestamp\"])\n",
    "df_test.set_index(df_test[\"timestamp\"], inplace=True)\n",
    "df_test.drop(columns=\"timestamp\", inplace=True)\n",
    "\n",
    "# Using user_provided time_col and freq\n",
    "timegpt_anomalies_df_1 = nixtla_client.detect_anomalies(df, time_col='timestamp', target_col='value', freq= 'M')\n",
    "# Infer time_col and freq from index\n",
    "timegpt_anomalies_df_2 = nixtla_client.detect_anomalies(df_test, time_col='timestamp', target_col='value')\n",
    "\n",
    "pd.testing.assert_frame_equal(\n",
    "    timegpt_anomalies_df_1,\n",
    "    timegpt_anomalies_df_2 \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "# test for showing the correct warning if X_df is missing but df has exogenous columns\n",
    "df = generate_series(n_series=2, min_length=5, max_length=20, n_static_features=3)\n",
    "missing_exogenous = df.columns.drop(['unique_id', 'ds', 'y']).tolist()\n",
    "expected_warning = (\n",
    "    f'`df` contains the following exogenous features: {missing_exogenous}, '\n",
    "    'but `X_df` was not provided. They will be ignored.'      \n",
    ")\n",
    "\n",
    "with warnings.catch_warnings(record=True) as w:\n",
    "    forecasts = nixtla_client.forecast(df, h=5)\n",
    "    assert any(expected_warning in str(warning.message) for warning in w)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| exporti\n",
    "def _forecast_wrapper(\n",
    "    df: pd.DataFrame,\n",
    "    client: NixtlaClient,\n",
    "    h: PositiveInt,\n",
    "    freq: Optional[str],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    level: Optional[List[Union[int, float]]],\n",
    "    quantiles: Optional[List[float]],\n",
    "    finetune_steps: NonNegativeInt,\n",
    "    finetune_loss: _Loss,\n",
    "    clean_ex_first: bool,\n",
    "    validate_api_key: bool,\n",
    "    add_history: bool,\n",
    "    date_features: Union[bool, List[Union[str, Callable]]],\n",
    "    date_features_to_one_hot: Union[bool, List[str]],\n",
    "    model: _Model,\n",
    "    num_partitions: Optional[PositiveInt],\n",
    ") -> pd.DataFrame:\n",
    "    if '_in_sample' in df:\n",
    "        in_sample_mask = df['_in_sample']\n",
    "        X_df = df.loc[~in_sample_mask].drop(columns=['_in_sample', target_col])\n",
    "        df = df.loc[in_sample_mask].drop(columns='_in_sample')\n",
    "    else:\n",
    "        X_df = None\n",
    "    return client.forecast(\n",
    "        df=df,\n",
    "        h=h,\n",
    "        freq=freq,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        X_df=X_df,\n",
    "        level=level,\n",
    "        quantiles=quantiles,\n",
    "        finetune_steps=finetune_steps,\n",
    "        finetune_loss=finetune_loss,\n",
    "        clean_ex_first=clean_ex_first,\n",
    "        validate_api_key=validate_api_key,\n",
    "        add_history=add_history,\n",
    "        date_features=date_features,\n",
    "        date_features_to_one_hot=date_features_to_one_hot,\n",
    "        model=model,\n",
    "        num_partitions=num_partitions,\n",
    "    )\n",
    "\n",
    "def _detect_anomalies_wrapper(\n",
    "    df: pd.DataFrame,\n",
    "    client: NixtlaClient,\n",
    "    freq: Optional[str],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    level: Union[int, float],\n",
    "    clean_ex_first: bool,\n",
    "    validate_api_key: bool,\n",
    "    date_features: Union[bool, List[str]],\n",
    "    date_features_to_one_hot: Union[bool, List[str]],\n",
    "    model: str,\n",
    "    num_partitions: Optional[PositiveInt],\n",
    ") -> pd.DataFrame:\n",
    "    return client.detect_anomalies(\n",
    "        df=df,\n",
    "        freq=freq,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        level=level,\n",
    "        clean_ex_first=clean_ex_first,\n",
    "        validate_api_key=validate_api_key,\n",
    "        date_features=date_features,\n",
    "        date_features_to_one_hot=date_features_to_one_hot,\n",
    "        model=model,\n",
    "        num_partitions=num_partitions,\n",
    "    )\n",
    "\n",
    "def _cross_validation_wrapper(\n",
    "    df: pd.DataFrame,\n",
    "    client: NixtlaClient,\n",
    "    h: PositiveInt,\n",
    "    freq: Optional[str],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    level: Optional[List[Union[int, float]]],\n",
    "    quantiles: Optional[List[float]],\n",
    "    validate_api_key: bool,\n",
    "    n_windows: PositiveInt,\n",
    "    step_size: Optional[PositiveInt],\n",
    "    finetune_steps: NonNegativeInt,\n",
    "    finetune_loss: str,\n",
    "    clean_ex_first: bool,\n",
    "    date_features: Union[bool, List[str]],\n",
    "    date_features_to_one_hot: Union[bool, List[str]],\n",
    "    model: str,\n",
    "    num_partitions: Optional[PositiveInt],\n",
    ") -> pd.DataFrame:\n",
    "    return client.cross_validation(\n",
    "        df=df,\n",
    "        h=h,\n",
    "        freq=freq,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        level=level,\n",
    "        quantiles=quantiles,\n",
    "        validate_api_key=validate_api_key,\n",
    "        n_windows=n_windows,\n",
    "        step_size=step_size,\n",
    "        finetune_steps=finetune_steps,\n",
    "        finetune_loss=finetune_loss,\n",
    "        clean_ex_first=clean_ex_first,\n",
    "        date_features=date_features,\n",
    "        date_features_to_one_hot=date_features_to_one_hot,\n",
    "        model=model,\n",
    "        num_partitions=num_partitions,\n",
    "    )\n",
    "\n",
    "def _get_schema(\n",
    "    df: 'fugue.AnyDataFrame',\n",
    "    method: str,\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    level: Optional[List[Union[int, float]]],\n",
    "    quantiles: Optional[List[float]],\n",
    ") -> 'triad.Schema':\n",
    "    import fugue.api as fa\n",
    "\n",
    "    base_cols = [id_col, time_col]\n",
    "    if method != 'forecast':\n",
    "        base_cols.append(target_col)\n",
    "    schema = fa.get_schema(df).extract(base_cols).copy()\n",
    "    schema.append('TimeGPT:double')\n",
    "    if method == 'detect_anomalies':\n",
    "        schema.append('anomaly:bool')\n",
    "    elif method == 'cross_validation':\n",
    "        schema.append(('cutoff', schema[time_col].type))\n",
    "    if level is not None and quantiles is not None:\n",
    "        raise ValueError(\"You should provide `level` or `quantiles` but not both.\")\n",
    "    if level is not None:\n",
    "        if not isinstance(level, list):\n",
    "            level = [level]\n",
    "        level = sorted(level)\n",
    "        schema.append(\",\".join(f\"TimeGPT-lo-{lv}:double\" for lv in reversed(level)))\n",
    "        schema.append(\",\".join(f\"TimeGPT-hi-{lv}:double\" for lv in level))\n",
    "    if quantiles is not None:\n",
    "        quantiles = sorted(quantiles)\n",
    "        q_cols = [f'TimeGPT-q-{int(q * 100)}:double' for q in quantiles]\n",
    "        schema.append(\",\".join(q_cols))\n",
    "    return schema\n",
    "\n",
    "def _distributed_setup(\n",
    "    df: 'fugue.AnyDataFrame',\n",
    "    method: str,\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    level: Optional[List[Union[int, float]]],\n",
    "    quantiles: Optional[List[float]],\n",
    "    num_partitions: Optional[int],\n",
    ") -> Tuple['triad.Schema', Dict[str, Any]]:\n",
    "    from fugue.execution import infer_execution_engine\n",
    "\n",
    "    if infer_execution_engine([df]) is None:\n",
    "        raise ValueError(\n",
    "            f'Could not infer execution engine for type {type(df).__name__}. '\n",
    "            'Expected a spark or dask DataFrame or a ray Dataset.'\n",
    "        )\n",
    "    schema = _get_schema(\n",
    "        df=df,\n",
    "        method=method,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        level=level,\n",
    "        quantiles=quantiles,\n",
    "    )\n",
    "    partition_config = dict(by=id_col, algo='coarse')\n",
    "    if num_partitions is not None:\n",
    "        partition_config['num'] = num_partitions\n",
    "    return schema, partition_config\n",
    "\n",
    "@patch\n",
    "def _distributed_forecast(\n",
    "    self: NixtlaClient,\n",
    "    df: 'fugue.AnyDataFrame',\n",
    "    h: PositiveInt,\n",
    "    freq: Optional[str],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    X_df: Optional['fugue.AnyDataFrame'],\n",
    "    level: Optional[List[Union[int, float]]],\n",
    "    quantiles: Optional[List[float]],\n",
    "    finetune_steps: NonNegativeInt,\n",
    "    finetune_loss: _Loss,\n",
    "    clean_ex_first: bool,\n",
    "    validate_api_key: bool,\n",
    "    add_history: bool,\n",
    "    date_features: Union[bool, List[Union[str, Callable]]],\n",
    "    date_features_to_one_hot: Union[bool, List[str]],\n",
    "    model: _Model,\n",
    "    num_partitions: Optional[int],\n",
    ") -> 'fugue.AnyDataFrame':\n",
    "    import fugue.api as fa\n",
    "\n",
    "    schema, partition_config = _distributed_setup(\n",
    "        df=df,\n",
    "        method='forecast',\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        level=level,\n",
    "        quantiles=quantiles,\n",
    "        num_partitions=num_partitions,\n",
    "    )\n",
    "    if X_df is not None:\n",
    "        def format_df(df: pd.DataFrame) -> pd.DataFrame:\n",
    "            return df.assign(_in_sample=True)\n",
    "\n",
    "        def format_X_df(\n",
    "            X_df: pd.DataFrame,\n",
    "            target_col: str,\n",
    "            df_cols: List[str],\n",
    "        ) -> pd.DataFrame:\n",
    "            return X_df.assign(**{'_in_sample': False, target_col: 0.0})[df_cols]\n",
    "\n",
    "        df = fa.transform(df, format_df, schema='*,_in_sample:bool')\n",
    "        X_df = fa.transform(\n",
    "            X_df,\n",
    "            format_X_df,\n",
    "            schema=fa.get_schema(df),\n",
    "            params={'target_col': target_col, 'df_cols': fa.get_column_names(df)},\n",
    "        )\n",
    "        df = fa.union(df, X_df)\n",
    "    result_df = fa.transform(\n",
    "        df,\n",
    "        using=_forecast_wrapper,\n",
    "        schema=schema,\n",
    "        params=dict(\n",
    "            client=self,\n",
    "            h=h,\n",
    "            freq=freq,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "            level=level,\n",
    "            quantiles=quantiles,\n",
    "            finetune_steps=finetune_steps,\n",
    "            finetune_loss=finetune_loss,\n",
    "            clean_ex_first=clean_ex_first,\n",
    "            validate_api_key=validate_api_key,\n",
    "            add_history=add_history,\n",
    "            date_features=date_features,\n",
    "            date_features_to_one_hot=date_features_to_one_hot,\n",
    "            model=model,\n",
    "            num_partitions=None,   \n",
    "        ),\n",
    "        partition=partition_config,\n",
    "        as_fugue=True,\n",
    "    )\n",
    "    return fa.get_native_as_df(result_df)\n",
    "\n",
    "@patch\n",
    "def _distributed_detect_anomalies(\n",
    "    self: NixtlaClient,\n",
    "    df: 'fugue.AnyDataFrame',\n",
    "    freq: Optional[str],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    level: Union[int, float],\n",
    "    clean_ex_first: bool,\n",
    "    validate_api_key: bool,\n",
    "    date_features: Union[bool, List[str]],\n",
    "    date_features_to_one_hot: Union[bool, List[str]],\n",
    "    model: str,\n",
    "    num_partitions: Optional[int],\n",
    "):\n",
    "    import fugue.api as fa\n",
    "\n",
    "    schema, partition_config = _distributed_setup(\n",
    "        df=df,\n",
    "        method='detect_anomalies',\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        level=level,\n",
    "        quantiles=None,\n",
    "        num_partitions=num_partitions,\n",
    "    )\n",
    "    result_df = fa.transform(\n",
    "        df,\n",
    "        using=_detect_anomalies_wrapper,\n",
    "        schema=schema,\n",
    "        params=dict(\n",
    "            client=self,\n",
    "            freq=freq,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "            level=level,\n",
    "            clean_ex_first=clean_ex_first,\n",
    "            validate_api_key=validate_api_key,\n",
    "            date_features=date_features,\n",
    "            date_features_to_one_hot=date_features_to_one_hot,\n",
    "            model=model,\n",
    "            num_partitions=None,   \n",
    "        ),\n",
    "        partition=partition_config,\n",
    "        as_fugue=True,\n",
    "    )\n",
    "    return fa.get_native_as_df(result_df)\n",
    "\n",
    "@patch\n",
    "def _distributed_cross_validation(\n",
    "    self: NixtlaClient,\n",
    "    df: 'fugue.AnyDataFrame',\n",
    "    h: PositiveInt,\n",
    "    freq: Optional[str],\n",
    "    id_col: str,\n",
    "    time_col: str,\n",
    "    target_col: str,\n",
    "    level: Optional[List[Union[int, float]]],\n",
    "    quantiles: Optional[List[float]],\n",
    "    validate_api_key: bool,\n",
    "    n_windows: PositiveInt,\n",
    "    step_size: Optional[PositiveInt],\n",
    "    finetune_steps: NonNegativeInt,\n",
    "    finetune_loss: _Loss,\n",
    "    clean_ex_first: bool,\n",
    "    date_features: Union[bool, List[Union[str, Callable]]],\n",
    "    date_features_to_one_hot: Union[bool, List[str]],\n",
    "    model: _Model,\n",
    "    num_partitions: Optional[int],\n",
    ") -> 'fugue.AnyDataFrame':\n",
    "    import fugue.api as fa\n",
    "\n",
    "    schema, partition_config = _distributed_setup(\n",
    "        df=df,\n",
    "        method='forecast',\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        level=level,\n",
    "        quantiles=quantiles,\n",
    "        num_partitions=num_partitions,\n",
    "    )\n",
    "    result_df = fa.transform(\n",
    "        df,\n",
    "        using=_cross_validation_wrapper,\n",
    "        schema=schema,\n",
    "        params=dict(\n",
    "            client=self,\n",
    "            h=h,\n",
    "            freq=freq,\n",
    "            id_col=id_col,\n",
    "            time_col=time_col,\n",
    "            target_col=target_col,\n",
    "            level=level,\n",
    "            quantiles=quantiles,\n",
    "            validate_api_key=validate_api_key,\n",
    "            n_windows=n_windows,\n",
    "            step_size=step_size,\n",
    "            finetune_steps=finetune_steps,\n",
    "            finetune_loss=finetune_loss,\n",
    "            clean_ex_first=clean_ex_first,\n",
    "            date_features=date_features,\n",
    "            date_features_to_one_hot=date_features_to_one_hot,\n",
    "            model=model,\n",
    "            num_partitions=None,   \n",
    "        ),\n",
    "        partition=partition_config,\n",
    "        as_fugue=True,\n",
    "    )\n",
    "    return fa.get_native_as_df(result_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| distributed\n",
    "import dask.dataframe as dd\n",
    "import fugue\n",
    "import fugue.api as fa\n",
    "import ray\n",
    "from dask.distributed import Client\n",
    "from pyspark.sql import SparkSession\n",
    "from ray.cluster_utils import Cluster"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| distributed\n",
    "def test_forecast(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    horizon: int = 12,\n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    **fcst_kwargs,\n",
    "):\n",
    "    fcst_df = nixtla_client.forecast(\n",
    "        df=df, \n",
    "        h=horizon,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        **fcst_kwargs,\n",
    "    )\n",
    "    fcst_df = fa.as_pandas(fcst_df)\n",
    "    test_eq(n_series * 12, len(fcst_df))\n",
    "    cols = fcst_df.columns.to_list()\n",
    "    exp_cols = [id_col, time_col, 'TimeGPT']\n",
    "    if 'level' in fcst_kwargs:\n",
    "        level = sorted(fcst_kwargs['level'])\n",
    "        exp_cols.extend([f'TimeGPT-lo-{lv}' for lv in reversed(level)])\n",
    "        exp_cols.extend([f'TimeGPT-hi-{lv}' for lv in level])\n",
    "    test_eq(cols, exp_cols)\n",
    "\n",
    "def test_forecast_diff_results_diff_models(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    horizon: int = 12, \n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    **fcst_kwargs,\n",
    "):\n",
    "    fcst_df = nixtla_client.forecast(\n",
    "        df=df, \n",
    "        h=horizon, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        model='timegpt-1',\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df = fa.as_pandas(fcst_df)\n",
    "    fcst_df_2 = nixtla_client.forecast(\n",
    "        df=df, \n",
    "        h=horizon, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        model='timegpt-1-long-horizon',\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df_2 = fa.as_pandas(fcst_df_2)\n",
    "    test_fail(\n",
    "        lambda: pd.testing.assert_frame_equal(\n",
    "            fcst_df.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "            fcst_df_2.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "        ),\n",
    "        contains='(column name=\"TimeGPT\") are different',\n",
    "    )\n",
    "\n",
    "def test_forecast_same_results_num_partitions(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    horizon: int = 12, \n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    **fcst_kwargs,\n",
    "):\n",
    "    fcst_df = nixtla_client.forecast(\n",
    "        df=df, \n",
    "        h=horizon, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df = fa.as_pandas(fcst_df)\n",
    "    fcst_df_2 = nixtla_client.forecast(\n",
    "        df=df, \n",
    "        h=horizon, \n",
    "        num_partitions=2,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df_2 = fa.as_pandas(fcst_df_2)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        fcst_df.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "        fcst_df_2.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "    )\n",
    "\n",
    "def test_cv_same_results_num_partitions(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    horizon: int = 12, \n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    **fcst_kwargs,\n",
    "):\n",
    "    fcst_df = nixtla_client.cross_validation(\n",
    "        df=df, \n",
    "        h=horizon, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df = fa.as_pandas(fcst_df)\n",
    "    fcst_df_2 = nixtla_client.cross_validation(\n",
    "        df=df, \n",
    "        h=horizon, \n",
    "        num_partitions=2,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df_2 = fa.as_pandas(fcst_df_2)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        fcst_df.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "        fcst_df_2.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "    )\n",
    "\n",
    "def test_forecast_dataframe(df: fugue.AnyDataFrame):\n",
    "    test_cv_same_results_num_partitions(df, n_windows=2, step_size=1)\n",
    "    test_cv_same_results_num_partitions(df, n_windows=3, step_size=None, horizon=1)\n",
    "    test_cv_same_results_num_partitions(df, model='timegpt-1-long-horizon', horizon=1)\n",
    "    test_forecast_diff_results_diff_models(df)\n",
    "    test_forecast(df, num_partitions=1)\n",
    "    test_forecast(df, level=[90, 80], num_partitions=1)\n",
    "    test_forecast_same_results_num_partitions(df)\n",
    "\n",
    "def test_forecast_dataframe_diff_cols(\n",
    "    df: fugue.AnyDataFrame,\n",
    "    id_col: str = 'id_col',\n",
    "    time_col: str = 'time_col',\n",
    "    target_col: str = 'target_col',\n",
    "):\n",
    "    test_forecast(df, id_col=id_col, time_col=time_col, target_col=target_col, num_partitions=1)\n",
    "    test_forecast(\n",
    "        df, id_col=id_col, time_col=time_col, target_col=target_col, level=[90, 80], num_partitions=1\n",
    "    )\n",
    "    test_forecast_same_results_num_partitions(\n",
    "        df, id_col=id_col, time_col=time_col, target_col=target_col\n",
    "    )\n",
    "\n",
    "def test_forecast_x(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    X_df: fugue.AnyDataFrame,\n",
    "    horizon: int = 24,\n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    target_col: str = 'y',\n",
    "    **fcst_kwargs,\n",
    "):\n",
    "    fcst_df = nixtla_client.forecast(\n",
    "        df=df, \n",
    "        X_df=X_df,\n",
    "        h=horizon,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        **fcst_kwargs,\n",
    "    )\n",
    "    fcst_df = fa.as_pandas(fcst_df)\n",
    "    n_series = fa.as_pandas(X_df)[id_col].nunique()\n",
    "    test_eq(n_series * horizon, len(fcst_df))\n",
    "    cols = fcst_df.columns.to_list()\n",
    "    exp_cols = [id_col, time_col, 'TimeGPT']\n",
    "    if 'level' in fcst_kwargs:\n",
    "        level = sorted(fcst_kwargs['level'])\n",
    "        exp_cols.extend([f'TimeGPT-lo-{lv}' for lv in reversed(level)])\n",
    "        exp_cols.extend([f'TimeGPT-hi-{lv}' for lv in level])\n",
    "    test_eq(cols, exp_cols)\n",
    "    fcst_df_2 = nixtla_client.forecast(\n",
    "        df=df,\n",
    "        h=horizon,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        **fcst_kwargs,\n",
    "    )\n",
    "    fcst_df_2 = fa.as_pandas(fcst_df_2)\n",
    "    equal_arrays = np.array_equal(\n",
    "        fcst_df.sort_values([id_col, time_col])['TimeGPT'].values,\n",
    "        fcst_df_2.sort_values([id_col, time_col])['TimeGPT'].values\n",
    "    )\n",
    "    assert not equal_arrays, 'Forecasts with and without ex vars are equal'\n",
    "\n",
    "def test_forecast_x_same_results_num_partitions(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    X_df: fugue.AnyDataFrame,\n",
    "    horizon: int = 24, \n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    target_col: str = 'y',\n",
    "    **fcst_kwargs,\n",
    "):\n",
    "    fcst_df = nixtla_client.forecast(\n",
    "        df=df, \n",
    "        X_df=X_df,\n",
    "        h=horizon, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df = fa.as_pandas(fcst_df)\n",
    "    fcst_df_2 = nixtla_client.forecast(\n",
    "        df=df,\n",
    "        h=horizon,\n",
    "        num_partitions=2,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        **fcst_kwargs\n",
    "    )\n",
    "    fcst_df_2 = fa.as_pandas(fcst_df_2)\n",
    "    equal_arrays = np.array_equal(\n",
    "        fcst_df.sort_values([id_col, time_col])['TimeGPT'].values,\n",
    "        fcst_df_2.sort_values([id_col, time_col])['TimeGPT'].values\n",
    "    )\n",
    "    assert not equal_arrays, 'Forecasts with and without ex vars are equal'\n",
    "\n",
    "def test_forecast_x_dataframe(df: fugue.AnyDataFrame, X_df: fugue.AnyDataFrame):\n",
    "    test_forecast_x(df, X_df, num_partitions=1)\n",
    "    test_forecast_x(df, X_df, level=[90, 80], num_partitions=1)\n",
    "    test_forecast_x_same_results_num_partitions(df, X_df)\n",
    "\n",
    "def test_forecast_x_dataframe_diff_cols(\n",
    "    df: fugue.AnyDataFrame,\n",
    "    X_df: fugue.AnyDataFrame,\n",
    "    id_col: str = 'id_col',\n",
    "    time_col: str = 'time_col',\n",
    "    target_col: str = 'target_col'\n",
    "):\n",
    "    test_forecast_x(\n",
    "        df, X_df, id_col=id_col, time_col=time_col, target_col=target_col, num_partitions=1\n",
    "    )\n",
    "    test_forecast_x(\n",
    "        df, X_df, id_col=id_col, time_col=time_col, target_col=target_col, level=[90, 80], num_partitions=1\n",
    "    )\n",
    "    test_forecast_x_same_results_num_partitions(\n",
    "        df, X_df, id_col=id_col, time_col=time_col, target_col=target_col\n",
    "    )\n",
    "\n",
    "def test_anomalies(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    target_col: str = 'y',\n",
    "    **anomalies_kwargs,\n",
    "):\n",
    "    anomalies_df = nixtla_client.detect_anomalies(\n",
    "        df=df, \n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        **anomalies_kwargs,\n",
    "    )\n",
    "    anomalies_df = fa.as_pandas(anomalies_df)\n",
    "    test_eq(fa.as_pandas(df)[id_col].unique(), anomalies_df[id_col].unique())\n",
    "    cols = anomalies_df.columns.to_list()\n",
    "    level = anomalies_kwargs.get('level', 99)\n",
    "    exp_cols = [\n",
    "        id_col,\n",
    "        time_col,\n",
    "        target_col,\n",
    "        'TimeGPT',\n",
    "        'anomaly',\n",
    "        f'TimeGPT-lo-{level}',\n",
    "        f'TimeGPT-hi-{level}',\n",
    "    ]\n",
    "    test_eq(cols, exp_cols)\n",
    "\n",
    "def test_anomalies_same_results_num_partitions(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    target_col: str = 'y',\n",
    "    **anomalies_kwargs,\n",
    "):\n",
    "    anomalies_df = nixtla_client.detect_anomalies(\n",
    "        df=df, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        **anomalies_kwargs\n",
    "    )\n",
    "    anomalies_df = fa.as_pandas(anomalies_df)\n",
    "    anomalies_df_2 = nixtla_client.detect_anomalies(\n",
    "        df=df, \n",
    "        num_partitions=2,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        **anomalies_kwargs\n",
    "    )\n",
    "    anomalies_df_2 = fa.as_pandas(anomalies_df_2)\n",
    "    pd.testing.assert_frame_equal(\n",
    "        anomalies_df.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "        anomalies_df_2.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "        atol=1e-5,\n",
    "    )\n",
    "\n",
    "def test_anomalies_diff_results_diff_models(\n",
    "    df: fugue.AnyDataFrame, \n",
    "    id_col: str = 'unique_id',\n",
    "    time_col: str = 'ds',\n",
    "    target_col: str = 'y',\n",
    "    **anomalies_kwargs,\n",
    "):\n",
    "    anomalies_df = nixtla_client.detect_anomalies(\n",
    "        df=df, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        model='timegpt-1',\n",
    "        **anomalies_kwargs\n",
    "    )\n",
    "    anomalies_df = fa.as_pandas(anomalies_df)\n",
    "    anomalies_df_2 = nixtla_client.detect_anomalies(\n",
    "        df=df, \n",
    "        num_partitions=1,\n",
    "        id_col=id_col,\n",
    "        time_col=time_col,\n",
    "        target_col=target_col,\n",
    "        model='timegpt-1-long-horizon',\n",
    "        **anomalies_kwargs\n",
    "    )\n",
    "    anomalies_df_2 = fa.as_pandas(anomalies_df_2)\n",
    "    test_fail(\n",
    "        lambda: pd.testing.assert_frame_equal(\n",
    "            anomalies_df.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "            anomalies_df_2.sort_values([id_col, time_col]).reset_index(drop=True),\n",
    "        ),\n",
    "        contains='(column name=\"TimeGPT\") are different',\n",
    "    )\n",
    "\n",
    "def test_anomalies_dataframe(df: fugue.AnyDataFrame):\n",
    "    test_anomalies(df, num_partitions=1)\n",
    "    test_anomalies(df, level=90, num_partitions=1)\n",
    "    test_anomalies_same_results_num_partitions(df)\n",
    "\n",
    "def test_anomalies_dataframe_diff_cols(\n",
    "    df: fugue.AnyDataFrame,\n",
    "    id_col: str = 'id_col',\n",
    "    time_col: str = 'time_col',\n",
    "    target_col: str = 'target_col',\n",
    "):\n",
    "    test_anomalies(df, id_col=id_col, time_col=time_col, target_col=target_col, num_partitions=1)\n",
    "    test_anomalies(df, id_col=id_col, time_col=time_col, target_col=target_col, level=90, num_partitions=1)\n",
    "    test_anomalies_same_results_num_partitions(df, id_col=id_col, time_col=time_col, target_col=target_col)\n",
    "    # @A: document behavior with exogenous variables in distributed environments.  \n",
    "    #test_anomalies_same_results_num_partitions(df, id_col=id_col, time_col=time_col, date_features=True, clean_ex_first=False)\n",
    "\n",
    "def test_quantiles(df: fugue.AnyDataFrame, id_col: str = 'id_col', time_col: str = 'time_col'):\n",
    "    test_qls = list(np.arange(0.1, 1, 0.1))\n",
    "    exp_q_cols = [f\"TimeGPT-q-{int(q * 100)}\" for q in test_qls]\n",
    "    def test_method_qls(method, **kwargs):\n",
    "        df_qls = method(\n",
    "            df=df, \n",
    "            h=12, \n",
    "            id_col=id_col,\n",
    "            time_col=time_col, \n",
    "            quantiles=test_qls,\n",
    "            **kwargs\n",
    "        )\n",
    "        df_qls = fa.as_pandas(df_qls)\n",
    "        assert all(col in df_qls.columns for col in exp_q_cols)\n",
    "        # test monotonicity of quantiles\n",
    "        df_qls.apply(lambda x: x.is_monotonic_increasing, axis=1).sum() == len(exp_q_cols)\n",
    "    test_method_qls(nixtla_client.forecast)\n",
    "    test_method_qls(nixtla_client.forecast, add_history=True)\n",
    "    test_method_qls(nixtla_client.cross_validation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| distributed\n",
    "n_series = 4\n",
    "horizon = 7\n",
    "\n",
    "series = generate_series(n_series, min_length=100)\n",
    "series['unique_id'] = series['unique_id'].astype(str)\n",
    "\n",
    "series_diff_cols = series.copy()\n",
    "renamer = {'unique_id': 'id_col', 'ds': 'time_col', 'y': 'target_col'}\n",
    "series_diff_cols = series_diff_cols.rename(columns=renamer)\n",
    "\n",
    "# data for exogenous tests\n",
    "df_x = pd.read_csv(\n",
    "    'https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/electricity-short-with-ex-vars.csv',\n",
    "    parse_dates=['ds'],\n",
    ")\n",
    "df_x = df_x.rename(columns=str.lower)\n",
    "future_ex_vars_df = pd.read_csv(\n",
    "    'https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/electricity-short-future-ex-vars.csv',\n",
    "    parse_dates=['ds'],\n",
    ")\n",
    "future_ex_vars_df = future_ex_vars_df.rename(columns=str.lower)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| distributed\n",
    "spark = SparkSession.builder.getOrCreate()\n",
    "spark_df = spark.createDataFrame(series).repartition(2)\n",
    "spark_diff_cols_df = spark.createDataFrame(series_diff_cols).repartition(2)\n",
    "\n",
    "test_quantiles(spark_df, id_col=\"unique_id\", time_col=\"ds\")\n",
    "\n",
    "test_forecast_dataframe(spark_df)\n",
    "test_forecast_dataframe_diff_cols(spark_diff_cols_df)\n",
    "test_anomalies_dataframe(spark_df)\n",
    "test_anomalies_dataframe_diff_cols(spark_diff_cols_df)\n",
    "# test exogenous variables\n",
    "spark_df_x = spark.createDataFrame(df_x).repartition(2)\n",
    "spark_future_ex_vars_df = spark.createDataFrame(future_ex_vars_df).repartition(2)\n",
    "test_forecast_x_dataframe(spark_df_x, spark_future_ex_vars_df)\n",
    "# test x different cols\n",
    "spark_df_x_diff_cols = spark.createDataFrame(df_x.rename(columns=renamer)).repartition(2)\n",
    "spark_future_ex_vars_df_diff_cols = spark.createDataFrame(\n",
    "    future_ex_vars_df.rename(columns=renamer)\n",
    ").repartition(2)\n",
    "test_forecast_x_dataframe_diff_cols(spark_df_x_diff_cols, spark_future_ex_vars_df_diff_cols)\n",
    "\n",
    "spark.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| distributed\n",
    "client = Client()\n",
    "dask_df = dd.from_pandas(series, npartitions=2)\n",
    "dask_diff_cols_df = dd.from_pandas(series_diff_cols, npartitions=2)\n",
    "\n",
    "test_quantiles(dask_df, id_col=\"unique_id\", time_col=\"ds\")\n",
    "\n",
    "\n",
    "test_forecast_dataframe(dask_df)\n",
    "test_forecast_dataframe_diff_cols(dask_diff_cols_df)\n",
    "test_anomalies_dataframe(dask_df)\n",
    "test_anomalies_dataframe_diff_cols(dask_diff_cols_df)\n",
    "\n",
    "# test exogenous variables\n",
    "dask_df_x = dd.from_pandas(df_x, npartitions=2)\n",
    "dask_future_ex_vars_df = dd.from_pandas(future_ex_vars_df, npartitions=2)\n",
    "test_forecast_x_dataframe(dask_df_x, dask_future_ex_vars_df)\n",
    "\n",
    "# test x different cols\n",
    "dask_df_x_diff_cols = dd.from_pandas(df_x.rename(columns=renamer), npartitions=2)\n",
    "dask_future_ex_vars_df_diff_cols = dd.from_pandas(future_ex_vars_df.rename(columns=renamer), npartitions=2)\n",
    "test_forecast_x_dataframe_diff_cols(dask_df_x_diff_cols, dask_future_ex_vars_df_diff_cols)\n",
    "\n",
    "client.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#| hide\n",
    "#| distributed\n",
    "ray_cluster = Cluster(\n",
    "    initialize_head=True,\n",
    "    head_node_args={\"num_cpus\": 2}\n",
    ")\n",
    "ray.init(address=ray_cluster.address, ignore_reinit_error=True)\n",
    "# add mock node to simulate a cluster\n",
    "mock_node = ray_cluster.add_node(num_cpus=2)\n",
    "\n",
    "ray_df = ray.data.from_pandas(series)\n",
    "ray_diff_cols_df = ray.data.from_pandas(series_diff_cols)\n",
    "\n",
    "test_quantiles(ray_df, id_col=\"unique_id\", time_col=\"ds\")\n",
    "\n",
    "test_forecast_dataframe(ray_df)\n",
    "test_forecast_dataframe_diff_cols(ray_diff_cols_df)\n",
    "test_anomalies_dataframe(ray_df)\n",
    "test_anomalies_dataframe_diff_cols(ray_diff_cols_df)\n",
    "\n",
    "# test exogenous variables\n",
    "ray_df_x = ray.data.from_pandas(df_x)\n",
    "ray_future_ex_vars_df = ray.data.from_pandas(future_ex_vars_df)\n",
    "test_forecast_x_dataframe(ray_df_x, ray_future_ex_vars_df)\n",
    "\n",
    "# test x different cols\n",
    "ray_df_x_diff_cols = ray.data.from_pandas(df_x.rename(columns=renamer))\n",
    "ray_future_ex_vars_df_diff_cols = ray.data.from_pandas(future_ex_vars_df.rename(columns=renamer))\n",
    "test_forecast_x_dataframe_diff_cols(ray_df_x_diff_cols, ray_future_ex_vars_df_diff_cols)\n",
    "\n",
    "ray.shutdown()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
